add fix for CUDA 10

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/14802

The definetion of THPStorage does not depend on any Real, its macro
defintion is unnecessary, refactor the code so that THPStorage is not macro
defined.

Reviewed By: ezyang

Differential Revision: D13340445

fbshipit-source-id: 343393d0a36c868b9a06eea2ad9b80f5e395e947

.clang-tidy

@@ -1,43 +1,31 @@
 ---
-# NOTE: there must be no spaces before the '-', so put the comma first.
+# NOTE there must be no spaces before the '-', so put the comma first.
 Checks: '
-  *
-  ,modernize-*
-  ,-cert-err58-cpp
-  ,-cert-err60-cpp
-  ,-clang-diagnostic-*
-  ,-cppcoreguidelines-owning-memory
+  -*
+  ,bugprone-*
+  ,-bugprone-macro-parentheses
+  ,-bugprone-forward-declaration-namespace
+  ,cppcoreguidelines-*
   ,-cppcoreguidelines-pro-bounds-array-to-pointer-decay
-  ,-cppcoreguidelines-pro-bounds-constant-array-index
   ,-cppcoreguidelines-pro-type-static-cast-downcast
+  ,-cppcoreguidelines-pro-bounds-pointer-arithmetic
+  ,-cppcoreguidelines-pro-bounds-constant-array-index
+  ,-cppcoreguidelines-pro-type-cstyle-cast
+  ,-cppcoreguidelines-pro-type-reinterpret-cast
   ,-cppcoreguidelines-pro-type-vararg
   ,-cppcoreguidelines-special-member-functions
-  ,-fuchsia-*
-  ,-google-build-using-namespace
-  ,-google-explicit-constructor
-  ,-google-readability-braces-around-statements
-  ,-google-readability-namespace-comments
-  ,-google-readability-todo
-  ,-google-runtime-references
-  ,-google-runtime-references
-  ,-hicpp-braces-around-statements
-  ,-hicpp-explicit-conversions
-  ,-hicpp-no-array-decay
-  ,-hicpp-special-member-functions
-  ,-hicpp-vararg
-  ,-llvm-header-guard
-  ,-llvm-namespace-comment
-  ,-misc-unused-parameters
-  ,-modernize-make-unique
+  ,-cppcoreguidelines-interfaces-global-init
+  ,-cppcoreguidelines-owning-memory
+  ,hicpp-signed-bitwise
+  ,hicpp-exception-baseclass
+  ,hicpp-avoid-goto
+  ,modernize-*
   ,-modernize-use-default-member-init
-  ,-performance-unnecessary-value-param
-  ,-readability-braces-around-statements
-  ,-readability-else-after-return
-  ,-readability-named-parameter
-  ,clang-analyzer-*
+  ,-modernize-return-braced-init-list
+  ,-modernize-use-auto
   '
-WarningsAsErrors: ''
-HeaderFilterRegex: 'torch/csrc/'
+WarningsAsErrors: '*'
+HeaderFilterRegex: 'torch/csrc/.*'
 AnalyzeTemporaryDtors: false
 CheckOptions:
 ...

.gitattributes

@@ -0,0 +1 @@
+*.bat	text eol=crlf

.github/ISSUE_TEMPLATE/bug-report.md

@@ -0,0 +1,49 @@
+---
+name: "\U0001F41B Bug Report"
+about: Submit a bug report to help us improve PyTorch
+
+---
+
+## 🐛 Bug
+
+<!-- A clear and concise description of what the bug is. -->
+
+## To Reproduce
+
+Steps to reproduce the behavior:
+
+1.
+1.
+1.
+
+<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->
+
+## Expected behavior
+
+<!-- A clear and concise description of what you expected to happen. -->
+
+## Environment
+
+Please copy and paste the output from our
+[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
+(or fill out the checklist below manually).
+
+You can get the script and run it with:
+```
+wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
+# For security purposes, please check the contents of collect_env.py before running it.
+python collect_env.py
+```
+
+ - PyTorch Version (e.g., 1.0):
+ - OS (e.g., Linux):
+ - How you installed PyTorch (`conda`, `pip`, source):
+ - Build command you used (if compiling from source):
+ - Python version:
+ - CUDA/cuDNN version:
+ - GPU models and configuration:
+ - Any other relevant information:
+
+## Additional context
+
+<!-- Add any other context about the problem here. -->

.github/ISSUE_TEMPLATE/documentation.md

@@ -0,0 +1,9 @@
+---
+name: "\U0001F4DA Documentation"
+about: Report an issue related to https://pytorch.org/docs
+
+---
+
+## 📚 Documentation
+
+<!-- A clear and concise description of what content in https://pytorch.org/docs is an issue. If this has to do with the general https://pytorch.org website, please file an issue at https://github.com/pytorch/pytorch.github.io/issues/new/choose instead. If this has to do with https://pytorch.org/tutorials, please file an issue at https://github.com/pytorch/tutorials/issues/new -->

.github/ISSUE_TEMPLATE/feature-request.md

@@ -0,0 +1,24 @@
+---
+name: "\U0001F680Feature Request"
+about: Submit a proposal/request for a new PyTorch feature
+
+---
+
+## 🚀 Feature
+<!-- A clear and concise description of the feature proposal -->
+
+## Motivation
+
+<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->
+
+## Pitch
+
+<!-- A clear and concise description of what you want to happen. -->
+
+## Alternatives
+
+<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->
+
+## Additional context
+
+<!-- Add any other context or screenshots about the feature request here. -->

.github/ISSUE_TEMPLATE/questions-help-support.md

@@ -0,0 +1,13 @@
+---
+name: "❓Questions/Help/Support"
+about: Do you need support? We have resources.
+
+---
+
+## ❓ Questions and Help
+
+### Please note that this issue tracker is not a help form and this issue will be closed.
+
+We have a set of [listed resources available on the website](https://pytorch.org/resources). Our primary means of support is our discussion forum:
+
+- [Discussion Forum](https://discuss.pytorch.org/)

.gitignore

@@ -25,13 +25,17 @@ aten/src/ATen/cuda/CUDAConfig.h
 build/
 dist/
 docs/src/**/*
+docs/cpp/build
+docs/cpp/source/api
 test/.coverage
 test/cpp/api/mnist
+test/custom_operator/model.pt
 test/data/gpu_tensors.pt
 test/data/legacy_modules.t7
 test/data/legacy_serialized.pt
 test/data/linear.pt
 test/htmlcov
+test/cpp_extensions/install/
 third_party/build/
 tools/shared/_utils_internal.py
 torch.egg-info/
@@ -40,6 +44,7 @@ torch/csrc/cudnn/cuDNN.cpp
 torch/csrc/generated
 torch/csrc/generic/TensorMethods.cpp
 torch/csrc/jit/generated/*
+torch/csrc/jit/fuser/config.h
 torch/csrc/nn/THCUNN.cpp
 torch/csrc/nn/THCUNN.cwrap
 torch/csrc/nn/THNN_generic.cpp
@@ -49,6 +54,7 @@ torch/csrc/nn/THNN.cpp
 torch/csrc/nn/THNN.cwrap
 torch/lib/*.a*
 torch/lib/*.dll*
+torch/lib/*.exe*
 torch/lib/*.dylib*
 torch/lib/*.h
 torch/lib/*.lib
@@ -60,6 +66,8 @@ torch/lib/pkgconfig
 torch/lib/protoc
 torch/lib/tmp_install
 torch/lib/torch_shm_manager
+torch/lib/python*
+torch/share/
 torch/version.py
 
 # IPython notebook checkpoints
@@ -140,10 +148,6 @@ docs/source/scripts/activation_images/
 # PyCharm files
 .idea
 
-# Visual Studio Code files
-.vscode
-.vs
-
 # OSX dir files
 .DS_Store
 
@@ -194,3 +198,26 @@ caffe2.egg-info
 
 # Atom/Watchman required file
 .watchmanconfig
+
+# Files generated by CLion
+cmake-build-debug
+
+# Files generated by ctags
+CTAGS
+tags
+TAGS
+
+# BEGIN NOT-CLEAN-FILES (setup.py handles this marker. Do not change.)
+#
+# Below files are not deleted by "setup.py clean".
+
+# Visual Studio Code files
+.vscode
+.vs
+
+# YouCompleteMe config file
+.ycm_extra_conf.py
+
+# Files generated when a patch is rejected
+*.orig
+*.rej

.gitmodules

@@ -1,9 +1,3 @@
-[submodule "third_party/catch"]
-	path = third_party/catch
-	url = https://github.com/catchorg/Catch2.git
-[submodule "third_party/nanopb"]
-	path = third_party/nanopb
-	url = https://github.com/nanopb/nanopb.git
 [submodule "third_party/pybind11"]
 	path = third_party/pybind11
 	url = https://github.com/pybind/pybind11.git
@@ -16,9 +10,6 @@
 [submodule "third_party/googletest"]
 	path = third_party/googletest
 	url = https://github.com/google/googletest.git
-[submodule "third_party/nervanagpu"]
-	path = third_party/nervanagpu
-	url = https://github.com/NervanaSystems/nervanagpu.git
 [submodule "third_party/benchmark"]
 	path = third_party/benchmark
 	url = https://github.com/google/benchmark.git
@@ -67,9 +58,6 @@
 [submodule "third_party/onnx"]
 	path = third_party/onnx
 	url = https://github.com/onnx/onnx.git
-[submodule "third_party/cereal"]
-	path = third_party/cereal
-	url = https://github.com/USCiLab/cereal
 [submodule "third_party/onnx-tensorrt"]
 	path = third_party/onnx-tensorrt
 	url = https://github.com/onnx/onnx-tensorrt
@@ -79,3 +67,18 @@
 [submodule "third_party/ideep"]
 	path = third_party/ideep
 	url = https://github.com/intel/ideep
+[submodule "third_party/nccl/nccl"]
+	path = third_party/nccl/nccl
+	url = https://github.com/NVIDIA/nccl
+[submodule "third_party/gemmlowp/gemmlowp"]
+	path = third_party/gemmlowp/gemmlowp
+	url = https://github.com/google/gemmlowp.git
+[submodule "third_party/QNNPACK"]
+	path = third_party/QNNPACK
+	url = https://github.com/pytorch/QNNPACK
+[submodule "third_party/neon2sse"]
+	path = third_party/neon2sse
+	url = https://github.com/intel/ARM_NEON_2_x86_SSE.git
+[submodule "third_party/fbgemm"]
+	path = third_party/fbgemm
+	url = https://github.com/pytorch/fbgemm

.jenkins/caffe2/build.sh

@@ -2,46 +2,57 @@
 
 set -ex
 
+pip install --user --no-cache-dir hypothesis==3.59.0
+
 # The INSTALL_PREFIX here must match up with test.sh
 INSTALL_PREFIX="/usr/local/caffe2"
 LOCAL_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
 ROOT_DIR=$(cd "$LOCAL_DIR"/../.. && pwd)
 CMAKE_ARGS=()
+SCCACHE="$(which sccache)"
 
+if [ "$(which gcc)" != "/root/sccache/gcc" ]; then
+  # Setup SCCACHE
+  ###############################################################################
+  # Setup sccache if SCCACHE_BUCKET is set
+  if [ -n "${SCCACHE_BUCKET}" ]; then
+    mkdir -p ./sccache
 
-# Setup SCCACHE
-###############################################################################
-# Setup sccache if SCCACHE_BUCKET is set
-if [ -n "${SCCACHE_BUCKET}" ]; then
-  mkdir -p ./sccache
+    SCCACHE="$(which sccache)"
+    if [ -z "${SCCACHE}" ]; then
+      echo "Unable to find sccache..."
+      exit 1
+    fi
 
-  SCCACHE="$(which sccache)"
-  if [ -z "${SCCACHE}" ]; then
-    echo "Unable to find sccache..."
-    exit 1
+    # Setup wrapper scripts
+    wrapped="cc c++ gcc g++ x86_64-linux-gnu-gcc"
+    if [[ "${BUILD_ENVIRONMENT}" == *-cuda* ]]; then
+        wrapped="$wrapped nvcc"
+    fi
+    for compiler in $wrapped; do
+      (
+        echo "#!/bin/sh"
+
+        # TODO: if/when sccache gains native support for an
+        # SCCACHE_DISABLE flag analogous to ccache's CCACHE_DISABLE,
+        # this can be removed. Alternatively, this can be removed when
+        # https://github.com/pytorch/pytorch/issues/13362 is fixed.
+        #
+        # NOTE: carefully quoted - we want `which compiler` to be
+        # resolved as we execute the script, but SCCACHE_DISABLE and
+        # $@ to be evaluated when we execute the script
+        echo 'test $SCCACHE_DISABLE && exec '"$(which $compiler)"' "$@"'
+
+        echo "exec $SCCACHE $(which $compiler) \"\$@\""
+      ) > "./sccache/$compiler"
+      chmod +x "./sccache/$compiler"
+    done
+
+    export CACHE_WRAPPER_DIR="$PWD/sccache"
+
+    # CMake must find these wrapper scripts
+    export PATH="$CACHE_WRAPPER_DIR:$PATH"
   fi
-
-  # Setup wrapper scripts
-  for compiler in cc c++ gcc g++ x86_64-linux-gnu-gcc; do
-    (
-      echo "#!/bin/sh"
-      echo "exec $SCCACHE $(which $compiler) \"\$@\""
-    ) > "./sccache/$compiler"
-    chmod +x "./sccache/$compiler"
-  done
-
-  if [[ "${BUILD_ENVIRONMENT}" == *-cuda* ]]; then
-    (
-      echo "#!/bin/sh"
-      echo "exec $SCCACHE $(which nvcc) \"\$@\""
-    ) > "./sccache/nvcc"
-    chmod +x "./sccache/nvcc"
-  fi
-
-  export CACHE_WRAPPER_DIR="$PWD/sccache"
-
-  # CMake must find these wrapper scripts
-  export PATH="$CACHE_WRAPPER_DIR:$PATH"
 fi
 
 # Setup ccache if configured to use it (and not sccache)
@@ -59,6 +70,15 @@ if [ -z "${SCCACHE}" ] && which ccache > /dev/null; then
   export PATH="$CACHE_WRAPPER_DIR:$PATH"
 fi
 
+# sccache will fail for CUDA builds if all cores are used for compiling
+if [ -z "$MAX_JOBS" ]; then
+  if [[ "${BUILD_ENVIRONMENT}" == *-cuda* ]] && [ -n "${SCCACHE}" ]; then
+    MAX_JOBS=`expr $(nproc) - 1`
+  else
+    MAX_JOBS=$(nproc)
+  fi
+fi
+
 report_compile_cache_stats() {
   if [[ -n "${SCCACHE}" ]]; then
     "$SCCACHE" --show-stats
@@ -79,7 +99,7 @@ fi
 
 
 ###############################################################################
-# Use special scripts for Android, conda, and setup builds
+# Use special scripts for Android and setup builds
 ###############################################################################
 if [[ "${BUILD_ENVIRONMENT}" == *-android* ]]; then
   export ANDROID_NDK=/opt/ndk
@@ -89,23 +109,6 @@ if [[ "${BUILD_ENVIRONMENT}" == *-android* ]]; then
   CMAKE_ARGS+=("-DUSE_ZSTD=ON")
   "${ROOT_DIR}/scripts/build_android.sh" ${CMAKE_ARGS[*]} "$@"
   exit 0
-elif [[ "${BUILD_ENVIRONMENT}" == conda* ]]; then
-  "${ROOT_DIR}/scripts/build_anaconda.sh" --skip-tests --install-locally "$@"
-  report_compile_cache_stats
-
-  # This build will be tested against onnx tests, which needs onnx installed.
-  # At this point the visible protbuf installation will be in conda, since one
-  # of Caffe2's dependencies uses conda, so the correct protobuf include
-  # headers are those in conda as well
-  # This path comes from install_anaconda.sh which installs Anaconda into the
-  # docker image
-  PROTOBUF_INCDIR=/opt/conda/include pip install -b /tmp/pip_install_onnx "file://${ROOT_DIR}/third_party/onnx#egg=onnx"
-  report_compile_cache_stats
-  exit 0
-elif [[ $BUILD_ENVIRONMENT == *setup* ]]; then
-  rm -rf $INSTALL_PREFIX && mkdir $INSTALL_PREFIX
-  PYTHONPATH=$INSTALL_PREFIX $PYTHON setup_caffe2.py develop --install-dir $INSTALL_PREFIX
-  exit 0
 fi
 
 
@@ -119,11 +122,6 @@ CMAKE_ARGS+=("-DUSE_OBSERVERS=ON")
 CMAKE_ARGS+=("-DUSE_ZSTD=ON")
 CMAKE_ARGS+=("-DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX}")
 
-if [[ $BUILD_ENVIRONMENT == *-aten-* ]]; then
-  if [[ CMAKE_ARGS != *USE_ATEN* ]] && [[ CMAKE_ARGS != *BUILD_ATEN* ]]; then
-    CMAKE_ARGS+=("-DBUILD_ATEN=ON")
-  fi
-fi
 if [[ $BUILD_ENVIRONMENT == *mkl* ]]; then
   CMAKE_ARGS+=("-DBLAS=MKL")
 fi
@@ -144,17 +142,19 @@ if [[ $BUILD_ENVIRONMENT == *cuda* ]]; then
   export PATH="/usr/local/cuda/bin:$PATH"
 fi
 if [[ $BUILD_ENVIRONMENT == *rocm* ]]; then
-  # TODO: This is patching the official FindHip to properly handly
-  # cmake generator expression. A PR is opened in the upstream repo here:
-  # https://github.com/ROCm-Developer-Tools/HIP/pull/516
-  # remove this hack once it's merged.
-  if [[ -f /opt/rocm/hip/cmake/FindHIP.cmake ]]; then
-    sudo sed -i 's/\ -I${dir}/\ $<$<BOOL:${dir}>:-I${dir}>/' /opt/rocm/hip/cmake/FindHIP.cmake
-  fi
+  # This is needed to enable ImageInput operator in resnet50_trainer
+  CMAKE_ARGS+=("-USE_OPENCV=ON")
+  # This is needed to read datasets from https://download.caffe2.ai/databases/resnet_trainer.zip
+  CMAKE_ARGS+=("-USE_LMDB=ON")
 
-  export LANG=C.UTF-8
-  export LC_ALL=C.UTF-8
-  export HCC_AMDGPU_TARGET=gfx900
+  ########## HIPIFY Caffe2 operators
+  ${PYTHON} "${ROOT_DIR}/tools/amd_build/build_amd.py"
+fi
+
+# building bundled nccl in this config triggers a bug in nvlink. For
+# more, see https://github.com/pytorch/pytorch/issues/14486
+if [[ "${BUILD_ENVIRONMENT}" == *-cuda8*-cudnn7* ]]; then
+    CMAKE_ARGS+=("-DUSE_SYSTEM_NCCL=ON")
 fi
 
 # Try to include Redis support for Linux builds
@@ -172,14 +172,6 @@ fi
 # Use a speciallized onnx namespace in CI to catch hardcoded onnx namespace
 CMAKE_ARGS+=("-DONNX_NAMESPACE=ONNX_NAMESPACE_FOR_C2_CI")
 
-if [[ -n "$INTEGRATED" ]]; then
-    # TODO: This is a temporary hack to work around the issue that both
-    # caffe2 and pytorch have libcaffe2.so and crossfire at runtime.
-    CMAKE_ARGS+=("-DBUILD_SHARED_LIBS=OFF")
-    CMAKE_ARGS+=("-DBUILD_CUSTOM_PROTOBUF=OFF")
-    CMAKE_ARGS+=("-DCAFFE2_LINK_LOCAL_PROTOBUF=OFF")
-fi
-
 # We test the presence of cmake3 (for platforms like Centos and Ubuntu 14.04)
 # and use that if so.
 if [[ -x "$(command -v cmake3)" ]]; then
@@ -187,38 +179,49 @@ if [[ -x "$(command -v cmake3)" ]]; then
 else
     CMAKE_BINARY=cmake
 fi
-# sccache will fail for CUDA builds if all cores are used for compiling
-if [[ "${BUILD_ENVIRONMENT}" == *-cuda* ]] && [ -n "${SCCACHE}" ]; then
-  MAX_JOBS=`expr $(nproc) - 1`
-else
-  MAX_JOBS=$(nproc)
-fi
-
 
 ###############################################################################
 # Configure and make
 ###############################################################################
-# Run cmake from ./build_caffe2 directory so it doesn't conflict with
-# standard PyTorch build directory. Eventually these won't need to
-# be separate.
-rm -rf build_caffe2
-mkdir build_caffe2
-cd ./build_caffe2
 
-# Configure
-${CMAKE_BINARY} "${ROOT_DIR}" ${CMAKE_ARGS[*]} "$@"
+if [[ -z "$INTEGRATED" ]]; then
+
+  # Run cmake from ./build_caffe2 directory so it doesn't conflict with
+  # standard PyTorch build directory. Eventually these won't need to
+  # be separate.
+  rm -rf build_caffe2
+  mkdir build_caffe2
+  cd ./build_caffe2
+
+  # Configure
+  ${CMAKE_BINARY} "${ROOT_DIR}" ${CMAKE_ARGS[*]} "$@"
+
+  # Build
+  if [ "$(uname)" == "Linux" ]; then
+    make "-j${MAX_JOBS}" install
+  else
+    echo "Don't know how to build on $(uname)"
+    exit 1
+  fi
 
-# Build
-if [ "$(uname)" == "Linux" ]; then
-  make "-j${MAX_JOBS}" install
 else
-  echo "Don't know how to build on $(uname)"
-  exit 1
+
+  # sccache will be stuck if  all cores are used for compiling
+  # see https://github.com/pytorch/pytorch/pull/7361
+  if [[ -n "${SCCACHE}" ]]; then
+    export MAX_JOBS=`expr $(nproc) - 1`
+  fi
+
+  USE_LEVELDB=1 USE_LMDB=1 USE_OPENCV=1 BUILD_BINARY=1 python setup.py install --user
+
+  # This is to save test binaries for testing
+  cp -r torch/lib/tmp_install $INSTALL_PREFIX
+
+  ls $INSTALL_PREFIX
+
+  report_compile_cache_stats
 fi
 
-report_compile_cache_stats
-
-
 ###############################################################################
 # Install ONNX
 ###############################################################################
@@ -228,17 +231,6 @@ pip install --user -b /tmp/pip_install_onnx "file://${ROOT_DIR}/third_party/onnx
 
 report_compile_cache_stats
 
-if [[ -n "$INTEGRATED" ]]; then
-  # sccache will be stuck if  all cores are used for compiling
-  # see https://github.com/pytorch/pytorch/pull/7361
-  if [[ -n "${SCCACHE}" ]]; then
-    export MAX_JOBS=`expr $(nproc) - 1`
-  fi
-  pip install --user -v -b /tmp/pip_install_torch "file://${ROOT_DIR}#egg=torch"
-fi
-
-report_compile_cache_stats
-
 # Symlink the caffe2 base python path into the system python path,
 # so that we can import caffe2 without having to change $PYTHONPATH.
 # Run in a subshell to contain environment set by /etc/os-release.
@@ -246,28 +238,30 @@ report_compile_cache_stats
 # This is only done when running on Jenkins!  We don't want to pollute
 # the user environment with Python symlinks and ld.so.conf.d hacks.
 #
-if [ -n "${JENKINS_URL}" ]; then
-  (
-    source /etc/os-release
+if [[ -z "$INTEGRATED" ]]; then
+  if [ -n "${JENKINS_URL}" ]; then
+    (
+      source /etc/os-release
 
-    function python_version() {
-      "$PYTHON" -c 'import sys; print("python%d.%d" % sys.version_info[0:2])'
-    }
+      function python_version() {
+        "$PYTHON" -c 'import sys; print("python%d.%d" % sys.version_info[0:2])'
+      }
 
-    # Debian/Ubuntu
-    if [[ "$ID_LIKE" == *debian* ]]; then
-      python_path="/usr/local/lib/$(python_version)/dist-packages"
-      sudo ln -sf "${INSTALL_PREFIX}/caffe2" "${python_path}"
-    fi
+      # Debian/Ubuntu
+      if [[ "$ID_LIKE" == *debian* ]]; then
+        python_path="/usr/local/lib/$(python_version)/dist-packages"
+        sudo ln -sf "${INSTALL_PREFIX}/caffe2" "${python_path}"
+      fi
 
-    # RHEL/CentOS
-    if [[ "$ID_LIKE" == *rhel* ]]; then
-      python_path="/usr/lib64/$(python_version)/site-packages/"
-      sudo ln -sf "${INSTALL_PREFIX}/caffe2" "${python_path}"
-    fi
+      # RHEL/CentOS
+      if [[ "$ID_LIKE" == *rhel* ]]; then
+        python_path="/usr/lib64/$(python_version)/site-packages/"
+        sudo ln -sf "${INSTALL_PREFIX}/caffe2" "${python_path}"
+      fi
 
-    # /etc/ld.so.conf.d is used on both Debian and RHEL
-    echo "${INSTALL_PREFIX}/lib" | sudo tee /etc/ld.so.conf.d/caffe2.conf
-    sudo ldconfig
-  )
+      # /etc/ld.so.conf.d is used on both Debian and RHEL
+      echo "${INSTALL_PREFIX}/lib" | sudo tee /etc/ld.so.conf.d/caffe2.conf
+      sudo ldconfig
+    )
+  fi
 fi

.jenkins/caffe2/test.sh

@@ -15,14 +15,6 @@ fi
 # The prefix must mirror the setting from build.sh
 INSTALL_PREFIX="/usr/local/caffe2"
 
-# Anaconda builds have a special install prefix and python
-if [[ "$BUILD_ENVIRONMENT" == conda* ]]; then
-  # This path comes from install_anaconda.sh which installs Anaconda into the
-  # docker image
-  PYTHON="/opt/conda/bin/python"
-  INSTALL_PREFIX="/opt/conda/"
-fi
-
 # Add the site-packages in the caffe2 install prefix to the PYTHONPATH
 SITE_DIR=$($PYTHON -c "from distutils import sysconfig; print(sysconfig.get_python_lib(prefix=''))")
 INSTALL_SITE_DIR="${INSTALL_PREFIX}/${SITE_DIR}"
@@ -34,11 +26,9 @@ if [[ "${BUILD_ENVIRONMENT}" == *-android* ]]; then
 fi
 
 # Set PYTHONPATH and LD_LIBRARY_PATH so that python can find the installed
-# Caffe2. This shouldn't be done on Anaconda, as Anaconda should handle this.
-if [[ "$BUILD_ENVIRONMENT" != conda* ]]; then
-  export PYTHONPATH="${PYTHONPATH}:$INSTALL_SITE_DIR"
-  export LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:${INSTALL_PREFIX}/lib"
-fi
+# Caffe2.
+export PYTHONPATH="${PYTHONPATH}:$INSTALL_SITE_DIR"
+export LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:${INSTALL_PREFIX}/lib"
 
 cd "$ROOT_DIR"
 
@@ -49,7 +39,7 @@ fi
 
 mkdir -p $TEST_DIR/{cpp,python}
 
-cd ${INSTALL_PREFIX}
+cd "${WORKSPACE}"
 
 # C++ tests
 echo "Running C++ tests.."
@@ -62,12 +52,26 @@ for test in $(find "${INSTALL_PREFIX}/test" -executable -type f); do
     */mkl_utils_test|*/aten/integer_divider_test)
       continue
       ;;
-    */aten/*)
-      # ATen uses test framework Catch2
-      "$test" -r=xml -o "${junit_reports_dir}/$(basename $test).xml"
-      ;;
-    *)
-      "$test" --gtest_output=xml:"$gtest_reports_dir/$(basename $test).xml"
+    */scalar_tensor_test|*/basic|*/native_test)
+	  if [[ "$BUILD_ENVIRONMENT" == *rocm* ]]; then
+		continue
+	  else
+	    "$test"
+	  fi
+	  ;;
+	*)
+      # Currently, we use a mixture of gtest (caffe2) and Catch2 (ATen). While
+      # planning to migrate to gtest as the common PyTorch c++ test suite, we
+      # currently do NOT use the xml test reporter, because Catch doesn't
+      # support multiple reporters
+      # c.f. https://github.com/catchorg/Catch2/blob/master/docs/release-notes.md#223
+      # which means that enabling XML output means you lose useful stdout
+      # output for Jenkins.  It's more important to have useful console
+      # output than it is to have XML output for Jenkins.
+      # Note: in the future, if we want to use xml test reporter once we switch
+      # to all gtest, one can simply do:
+      # "$test" --gtest_output=xml:"$gtest_reports_dir/$(basename $test).xml"
+      "$test"
       ;;
   esac
 done
@@ -83,33 +87,39 @@ if [[ "$BUILD_ENVIRONMENT" == *-cuda* ]]; then
   EXTRA_TESTS+=("$CAFFE2_PYPATH/contrib/nccl")
 fi
 
-conda_ignore_test=()
-if [[ $BUILD_ENVIRONMENT == conda* ]]; then
-  # These tests both assume Caffe2 was built with leveldb, which is not the case
-  conda_ignore_test+=("--ignore $CAFFE2_PYPATH/python/dataio_test.py")
-  conda_ignore_test+=("--ignore $CAFFE2_PYPATH/python/operator_test/checkpoint_test.py")
+rocm_ignore_test=()
+if [[ $BUILD_ENVIRONMENT == *-rocm* ]]; then
+  # Currently these tests are failing on ROCM platform:
+
+  # Unknown reasons, need to debug
+  rocm_ignore_test+=("--ignore $CAFFE2_PYPATH/python/operator_test/arg_ops_test.py")
+  rocm_ignore_test+=("--ignore $CAFFE2_PYPATH/python/operator_test/piecewise_linear_transform_test.py")
+  rocm_ignore_test+=("--ignore $CAFFE2_PYPATH/python/operator_test/softmax_ops_test.py")
+  rocm_ignore_test+=("--ignore $CAFFE2_PYPATH/python/operator_test/unique_ops_test.py")
 fi
 
+# Python tests
+# NB: Warnings are disabled because they make it harder to see what
+# the actual erroring test is
+echo "Running Python tests.."
+pip install --user pytest-sugar
+"$PYTHON" \
+  -m pytest \
+  -x \
+  -v \
+  --disable-warnings \
+  --junit-xml="$TEST_DIR/python/result.xml" \
+  --ignore "$CAFFE2_PYPATH/python/test/executor_test.py" \
+  --ignore "$CAFFE2_PYPATH/python/operator_test/matmul_op_test.py" \
+  --ignore "$CAFFE2_PYPATH/python/operator_test/pack_ops_test.py" \
+  --ignore "$CAFFE2_PYPATH/python/mkl/mkl_sbn_speed_test.py" \
+  ${rocm_ignore_test[@]} \
+  "$CAFFE2_PYPATH/python" \
+  "${EXTRA_TESTS[@]}"
 
-# TODO: re-enable this for rocm CI jobs once we have more rocm workers
-if [[ $BUILD_ENVIRONMENT != *rocm* ]]; then
-  # Python tests
-  echo "Running Python tests.."
-  "$PYTHON" \
-    -m pytest \
-    -x \
-    -v \
-    --junit-xml="$TEST_DIR/python/result.xml" \
-    --ignore "$CAFFE2_PYPATH/python/test/executor_test.py" \
-    --ignore "$CAFFE2_PYPATH/python/operator_test/matmul_op_test.py" \
-    --ignore "$CAFFE2_PYPATH/python/operator_test/pack_ops_test.py" \
-    --ignore "$CAFFE2_PYPATH/python/mkl/mkl_sbn_speed_test.py" \
-    ${conda_ignore_test[@]} \
-    "$CAFFE2_PYPATH/python" \
-    "${EXTRA_TESTS[@]}"
-fi
+cd ${INSTALL_PREFIX}
 
 if [[ -n "$INTEGRATED" ]]; then
-  pip install --user pytest-xdist torchvision
-  "$ROOT_DIR/scripts/onnx/test.sh" -p
+  pip install --user torchvision
+  "$ROOT_DIR/scripts/onnx/test.sh"
 fi

.jenkins/pytorch/build-asan.sh

@@ -14,8 +14,18 @@ clang --version
 # symbolize=1: Gives us much better errors when things go wrong
 export ASAN_OPTIONS=detect_leaks=0:symbolize=1
 
+# FIXME: Remove the hardcoded "-pthread" option.
+# With asan build, the cmake thread CMAKE_HAVE_LIBC_CREATE[1] checking will
+# succeed because "pthread_create" is in libasan.so. However, libasan doesn't
+# have the full pthread implementation. Other advanced pthread functions doesn't
+# exist in libasan.so[2]. If we need some pthread advanced functions, we still
+# need to link the pthread library.
+# [1] https://github.com/Kitware/CMake/blob/8cabaaf054a16ea9c8332ce8e9291bd026b38c62/Modules/FindThreads.cmake#L135
+# [2] https://wiki.gentoo.org/wiki/AddressSanitizer/Problems
+#
 # TODO: Make the ASAN flags a more unified env var
 CC="clang" CXX="clang++" LDSHARED="clang --shared" \
-  CFLAGS="-fsanitize=address -fsanitize=undefined -fno-sanitize-recover=all -shared-libasan" \
-  NO_CUDA=1 DEBUG=1 \
+  CFLAGS="-fsanitize=address -fsanitize=undefined -fno-sanitize-recover=all -shared-libasan -pthread" \
+  CXX_FLAGS="-pthread" \
+  NO_CUDA=1 USE_MKLDNN=0 \
   python setup.py install

.jenkins/pytorch/build.sh

@@ -1,15 +1,5 @@
 #!/bin/bash
 
-if [[ "$BUILD_ENVIRONMENT" == "pytorch-linux-xenial-py3-clang5-asan" ]]; then
-  exec "$(dirname "${BASH_SOURCE[0]}")/build-asan.sh" $*
-fi
-
-# TODO: move this to Docker
-# TODO: add both NCCL and MPI in CI test by fixing these test first
-# sudo apt-get update
-# sudo apt-get install libnccl-dev libnccl2
-# sudo apt-get install openmpi-bin libopenmpi-dev
-
 # Required environment variable: $BUILD_ENVIRONMENT
 # (This is set by default in the Docker images we build, so you don't
 # need to set it yourself.
@@ -17,6 +7,33 @@ fi
 COMPACT_JOB_NAME="${BUILD_ENVIRONMENT}-build"
 source "$(dirname "${BASH_SOURCE[0]}")/common.sh"
 
+# For distributed, four environmental configs:
+# (1) build with only NCCL
+# (2) build with NCCL and MPI
+# (3) build with only MPI
+# (4) build with neither
+if [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda9-* ]]; then
+  # TODO: move this to Docker
+  sudo apt-get -qq update
+  sudo apt-get -qq install --allow-downgrades --allow-change-held-packages libnccl-dev=2.2.13-1+cuda9.0 libnccl2=2.2.13-1+cuda9.0
+fi
+
+if [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda9*gcc7* ]] || [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda8-* ]] || [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda9-cudnn7-py2* ]] || [[ "$BUILD_ENVIRONMENT" == *-trusty-py2.7.9* ]]; then
+  # TODO: move this to Docker
+  sudo apt-get -qq update
+  if [[ "$BUILD_ENVIRONMENT" == *-trusty-py2.7.9* ]]; then
+    sudo apt-get -qq install openmpi-bin libopenmpi-dev
+  else
+    sudo apt-get -qq install --allow-downgrades --allow-change-held-packages openmpi-bin libopenmpi-dev
+  fi
+  sudo apt-get -qq install --no-install-recommends openssh-client openssh-server
+  sudo mkdir -p /var/run/sshd
+fi
+
+if [[ "$BUILD_ENVIRONMENT" == "pytorch-linux-xenial-py3-clang5-asan" ]]; then
+  exec "$(dirname "${BASH_SOURCE[0]}")/build-asan.sh" $*
+fi
+
 echo "Python version:"
 python --version
 
@@ -27,66 +44,105 @@ echo "CMake version:"
 cmake --version
 
 # TODO: Don't run this...
-pip install -r requirements.txt || true
+pip install -q -r requirements.txt || true
 
 if [[ "$BUILD_ENVIRONMENT" == *rocm* ]]; then
-  export HCC_AMDGPU_TARGET=gfx900
-  export LANG=C.UTF-8
-  export LC_ALL=C.UTF-8
+  # When hcc runs out of memory, it silently exits without stopping
+  # the build process, leaving undefined symbols in the shared lib
+  # which will cause undefined symbol errors when later running
+  # tests. Setting MAX_JOBS to smaller number to make CI less flaky.
+  export MAX_JOBS=4
 
-  sudo chown -R jenkins:jenkins /usr/local
-  rm -rf "$(dirname "${BASH_SOURCE[0]}")/../../../pytorch_amd/" || true
-  python "$(dirname "${BASH_SOURCE[0]}")/../../tools/amd_build/build_pytorch_amd.py"
-  USE_ROCM=1 python setup.py install
-  exit
+  # ROCm CI is using Caffe2 docker images, which needs these wrapper
+  # scripts to correctly use sccache.
+  if [ -n "${SCCACHE_BUCKET}" ]; then
+    mkdir -p ./sccache
+
+    SCCACHE="$(which sccache)"
+    if [ -z "${SCCACHE}" ]; then
+      echo "Unable to find sccache..."
+      exit 1
+    fi
+
+    # Setup wrapper scripts
+    for compiler in cc c++ gcc g++ x86_64-linux-gnu-gcc; do
+      (
+        echo "#!/bin/sh"
+        echo "exec $SCCACHE $(which $compiler) \"\$@\""
+      ) > "./sccache/$compiler"
+      chmod +x "./sccache/$compiler"
+    done
+
+    export CACHE_WRAPPER_DIR="$PWD/sccache"
+
+    # CMake must find these wrapper scripts
+    export PATH="$CACHE_WRAPPER_DIR:$PATH"
+  fi
+
+  python tools/amd_build/build_amd.py
+  # OPENCV is needed to enable ImageInput operator in caffe2 resnet5_trainer
+  # LMDB is needed to read datasets from https://download.caffe2.ai/databases/resnet_trainer.zip
+  USE_ROCM=1 USE_LMDB=1 USE_OPENCV=1 python setup.py install --user
+  exit 0
 fi
 
 # TODO: Don't install this here
 if ! which conda; then
-  pip install mkl mkl-devel
+  pip install -q mkl mkl-devel
+  if [[ "$BUILD_ENVIRONMENT" == *trusty-py3.6-gcc7.2* ]] || [[ "$BUILD_ENVIRONMENT" == *trusty-py3.6-gcc4.8* ]]; then
+    export USE_MKLDNN=1
+  else
+    export USE_MKLDNN=0
+  fi
 fi
 
 # sccache will fail for CUDA builds if all cores are used for compiling
 # gcc 7 with sccache seems to have intermittent OOM issue if all cores are used
-if ([[ "$BUILD_ENVIRONMENT" == *cuda* ]] || [[ "$BUILD_ENVIRONMENT" == *gcc7* ]]) && which sccache > /dev/null; then
-  export MAX_JOBS=`expr $(nproc) - 1`
+if [ -z "$MAX_JOBS" ]; then
+  if ([[ "$BUILD_ENVIRONMENT" == *cuda* ]] || [[ "$BUILD_ENVIRONMENT" == *gcc7* ]]) && which sccache > /dev/null; then
+    export MAX_JOBS=`expr $(nproc) - 1`
+  fi
 fi
 
 # Target only our CI GPU machine's CUDA arch to speed up the build
-export TORCH_CUDA_ARCH_LIST=5.2
+export TORCH_CUDA_ARCH_LIST="5.2"
+
+if [[ "$BUILD_ENVIRONMENT" == *ppc64le* ]]; then
+  export TORCH_CUDA_ARCH_LIST="6.0"
+fi
 
 if [[ "$BUILD_ENVIRONMENT" == *trusty-py3.6-gcc5.4* ]]; then
   export DEBUG=1
 fi
 
-WERROR=1 python setup.py install
+# ppc64le build fails when WERROR=1
+# set only when building other architectures
+# only use for "python setup.py install" line
+if [[ "$BUILD_ENVIRONMENT" != *ppc64le* ]]; then
+  WERROR=1 python setup.py install
+elif [[ "$BUILD_ENVIRONMENT" == *ppc64le* ]]; then
+  python setup.py install
+fi
+
 
 # Add the test binaries so that they won't be git clean'ed away
 git add -f build/bin
 
-# Testing ATen install
-if [[ "$BUILD_ENVIRONMENT" != *cuda* ]]; then
-  echo "Testing ATen install"
-  time tools/test_aten_install.sh
-fi
-
-# Test C FFI plugins
-# cffi install doesn't work for Python 3.7
-if [[ "$BUILD_ENVIRONMENT" != *pynightly* ]]; then
-  # TODO: Don't run this here
-  pip install cffi
-  git clone https://github.com/pytorch/extension-ffi.git
-  pushd extension-ffi/script
-  python build.py
-  popd
-fi
-
 # Test documentation build
 if [[ "$BUILD_ENVIRONMENT" == *xenial-cuda8-cudnn6-py3* ]]; then
   pushd docs
   # TODO: Don't run this here
-  pip install -r requirements.txt || true
-  make html
+  pip install -q -r requirements.txt || true
+  LC_ALL=C make html
+  popd
+fi
+
+# Test standalone c10 build
+if [[ "$BUILD_ENVIRONMENT" == *xenial-cuda8-cudnn6-py3* ]]; then
+  mkdir -p c10/build
+  pushd c10/build
+  cmake ..
+  make -j
   popd
 fi
 
@@ -95,5 +151,19 @@ 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_all.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
+
+  # Build custom operator tests.
+  CUSTOM_OP_BUILD="$PWD/../custom-op-build"
+  CUSTOM_OP_TEST="$PWD/test/custom_operator"
+  SITE_PACKAGES="$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')"
+  mkdir "$CUSTOM_OP_BUILD"
+  pushd "$CUSTOM_OP_BUILD"
+  CMAKE_PREFIX_PATH="$SITE_PACKAGES/torch" cmake "$CUSTOM_OP_TEST"
+  make VERBOSE=1
+  popd
 fi

.jenkins/pytorch/enabled-configs.txt

@@ -14,6 +14,8 @@ pytorch-linux-xenial-cuda9-cudnn7-py3-build
 pytorch-linux-xenial-cuda9-cudnn7-py3-test
 pytorch-linux-xenial-cuda9.2-cudnn7-py3-gcc7-build
 pytorch-linux-xenial-cuda9.2-cudnn7-py3-gcc7-test
+pytorch-linux-xenial-cuda10-cudnn7-py3-gcc7-build
+pytorch-linux-xenial-cuda10-cudnn7-py3-gcc7-test
 pytorch-linux-xenial-py3-clang5-asan-build
 pytorch-linux-xenial-py3-clang5-asan-test
 pytorch-linux-trusty-py2.7.9-build
@@ -40,4 +42,10 @@ pytorch-macos-10.13-cuda9.2-cudnn7-py3-build
 pytorch-docker-build-test
 short-perf-test-cpu
 short-perf-test-gpu
-py2-clang3.8-rocmnightly-ubuntu16.04-build
+py2-clang7-rocmdeb-ubuntu16.04-build
+py2-clang7-rocmdeb-ubuntu16.04-test
+py2-devtoolset7-rocmrpm-centos7.5-build
+pytorch-ppc64le-cuda9.2-cudnn7-py3-build
+pytorch-ppc64le-cuda9.2-cudnn7-py3-test
+pytorch-ppc64le-cuda9.1-cudnn7-py3-build
+pytorch-ppc64le-cuda9.1-cudnn7-py3-test

.jenkins/pytorch/macos-build.sh

@@ -29,11 +29,15 @@ if [[ "${JOB_BASE_NAME}" == *cuda9.2* ]]; then
   export CUDA_HOME=/Developer/NVIDIA/CUDA-${CUDA_VERSION}
   export NO_CUDA=0
 
-  # Eigen gives "explicit specialization of class must precede its first use" error
-  # when compiling with Xcode 9.1 toolchain, so we have to use Xcode 8.2 toolchain instead.
-  export DEVELOPER_DIR=/Library/Developer/CommandLineTools
+  if [ -z "${IN_CIRCLECI}" ]; then
+    # Eigen gives "explicit specialization of class must precede its first use" error
+    # when compiling with Xcode 9.1 toolchain, so we have to use Xcode 8.2 toolchain instead.
+    export DEVELOPER_DIR=/Library/Developer/CommandLineTools
+  fi
 else
-  export DEVELOPER_DIR=/Applications/Xcode9.app/Contents/Developer
+  if [ -z "${IN_CIRCLECI}" ]; then
+    export DEVELOPER_DIR=/Applications/Xcode9.app/Contents/Developer
+  fi
 fi
 
 export MACOSX_DEPLOYMENT_TARGET=10.9
@@ -62,5 +66,7 @@ export IMAGE_COMMIT_TAG=${BUILD_ENVIRONMENT}-${IMAGE_COMMIT_ID}
 python setup.py install
 
 # Upload torch binaries when the build job is finished
-7z a ${IMAGE_COMMIT_TAG}.7z ${PYTORCH_ENV_DIR}/miniconda3/lib/python3.6/site-packages/torch*
-aws s3 cp ${IMAGE_COMMIT_TAG}.7z s3://ossci-macos-build/pytorch/${IMAGE_COMMIT_TAG}.7z --acl public-read
+if [ -z "${IN_CIRCLECI}" ]; then
+  7z a ${IMAGE_COMMIT_TAG}.7z ${PYTORCH_ENV_DIR}/miniconda3/lib/python3.6/site-packages/torch*
+  aws s3 cp ${IMAGE_COMMIT_TAG}.7z s3://ossci-macos-build/pytorch/${IMAGE_COMMIT_TAG}.7z --acl public-read
+fi

.jenkins/pytorch/macos-test.sh

@@ -15,19 +15,24 @@ if [ ! -d "${PYTORCH_ENV_DIR}/miniconda3" ]; then
 fi
 export PATH="${PYTORCH_ENV_DIR}/miniconda3/bin:$PATH"
 source ${PYTORCH_ENV_DIR}/miniconda3/bin/activate
-conda install -y mkl mkl-include numpy pyyaml setuptools cmake cffi ninja
-rm -rf ${PYTORCH_ENV_DIR}/miniconda3/lib/python3.6/site-packages/torch*
+conda install -y mkl mkl-include numpy pyyaml setuptools cmake cffi ninja six
+pip install hypothesis
+if [ -z "${IN_CIRCLECI}" ]; then
+  rm -rf ${PYTORCH_ENV_DIR}/miniconda3/lib/python3.6/site-packages/torch*
+fi
 
 git submodule update --init --recursive
 export CMAKE_PREFIX_PATH=${PYTORCH_ENV_DIR}/miniconda3/
 
 # Test PyTorch
-if [[ "${JOB_BASE_NAME}" == *cuda9.2* ]]; then
-  # Eigen gives "explicit specialization of class must precede its first use" error
-  # when compiling with Xcode 9.1 toolchain, so we have to use Xcode 8.2 toolchain instead.
-  export DEVELOPER_DIR=/Library/Developer/CommandLineTools
-else
-  export DEVELOPER_DIR=/Applications/Xcode9.app/Contents/Developer
+if [ -z "${IN_CIRCLECI}" ]; then
+  if [[ "${JOB_BASE_NAME}" == *cuda9.2* ]]; then
+    # Eigen gives "explicit specialization of class must precede its first use" error
+    # when compiling with Xcode 9.1 toolchain, so we have to use Xcode 8.2 toolchain instead.
+    export DEVELOPER_DIR=/Library/Developer/CommandLineTools
+  else
+    export DEVELOPER_DIR=/Applications/Xcode9.app/Contents/Developer
+  fi
 fi
 export MACOSX_DEPLOYMENT_TARGET=10.9
 export CXX=clang++
@@ -38,9 +43,11 @@ export MAX_JOBS=2
 export IMAGE_COMMIT_TAG=${BUILD_ENVIRONMENT}-${IMAGE_COMMIT_ID}
 
 # Download torch binaries in the test jobs
-rm -rf ${PYTORCH_ENV_DIR}/miniconda3/lib/python3.6/site-packages/torch*
-aws s3 cp s3://ossci-macos-build/pytorch/${IMAGE_COMMIT_TAG}.7z ${IMAGE_COMMIT_TAG}.7z
-7z x ${IMAGE_COMMIT_TAG}.7z -o"${PYTORCH_ENV_DIR}/miniconda3/lib/python3.6/site-packages"
+if [ -z "${IN_CIRCLECI}" ]; then
+  rm -rf ${PYTORCH_ENV_DIR}/miniconda3/lib/python3.6/site-packages/torch*
+  aws s3 cp s3://ossci-macos-build/pytorch/${IMAGE_COMMIT_TAG}.7z ${IMAGE_COMMIT_TAG}.7z
+  7z x ${IMAGE_COMMIT_TAG}.7z -o"${PYTORCH_ENV_DIR}/miniconda3/lib/python3.6/site-packages"
+fi
 
 test_python_all() {
   echo "Ninja version: $(ninja --version)"
@@ -56,8 +63,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_all.sh "$CPP_BUILD"
+  mkdir -p $CPP_BUILD/caffe2
+
+  BUILD_LIBTORCH_PY=$PWD/tools/build_libtorch.py
+  pushd $CPP_BUILD/caffe2
+  VERBOSE=1 DEBUG=1 python $BUILD_LIBTORCH_PY
+  popd
 
   python tools/download_mnist.py --quiet -d test/cpp/api/mnist
 
@@ -65,16 +76,38 @@ test_cpp_api() {
   # without these paths being set
   export DYLD_LIBRARY_PATH="$DYLD_LIBRARY_PATH:$PWD/miniconda3/lib"
   export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:$PWD/miniconda3/lib"
-  "$CPP_BUILD"/libtorch/bin/test_api
+  "$CPP_BUILD"/caffe2/bin/test_api
 }
 
+test_custom_script_ops() {
+  echo "Testing custom script operators"
+  pushd test/custom_operator
+  # Build the custom operator library.
+  rm -rf build && mkdir build
+  pushd build
+  SITE_PACKAGES="$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')"
+  CMAKE_PREFIX_PATH="$SITE_PACKAGES/torch" cmake ..
+  make VERBOSE=1
+  popd
+
+  # Run tests Python-side and export a script module.
+  python test_custom_ops.py -v
+  python model.py --export-script-module=model.pt
+  # Run tests C++-side and load the exported script module.
+  build/test_custom_ops ./model.pt
+  popd
+}
+
+
 if [ -z "${JOB_BASE_NAME}" ] || [[ "${JOB_BASE_NAME}" == *-test ]]; then
   test_python_all
   test_cpp_api
+  test_custom_script_ops
 else
   if [[ "${JOB_BASE_NAME}" == *-test1 ]]; then
     test_python_all
   elif [[ "${JOB_BASE_NAME}" == *-test2 ]]; then
     test_cpp_api
+    test_custom_script_ops
   fi
 fi

.jenkins/pytorch/multigpu-test.sh

@@ -8,4 +8,21 @@ COMPACT_JOB_NAME="${BUILD_ENVIRONMENT}-multigpu-test"
 source "$(dirname "${BASH_SOURCE[0]}")/common.sh"
 
 echo "Testing pytorch (distributed only)"
+
+if [ -n "${IN_CIRCLECI}" ]; then
+  if [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda9-* ]]; then
+    # TODO: move this to Docker
+    sudo apt-get update
+    sudo apt-get install -y --allow-downgrades --allow-change-held-packages libnccl-dev=2.2.13-1+cuda9.0 libnccl2=2.2.13-1+cuda9.0
+  fi
+
+  if [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda8-* ]] || [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda9-cudnn7-py2* ]]; then
+    # TODO: move this to Docker
+    sudo apt-get update
+    sudo apt-get install -y --allow-downgrades --allow-change-held-packages openmpi-bin libopenmpi-dev
+    sudo apt-get install -y --no-install-recommends openssh-client openssh-server
+    sudo mkdir -p /var/run/sshd
+  fi
+fi
+
 time python test/run_test.py --verbose -i distributed

.jenkins/pytorch/perf_test/compare_with_baseline.py

@@ -1,5 +1,6 @@
 import sys
 import json
+import math
 import numpy
 import argparse
 
@@ -35,14 +36,25 @@ else:
 print("population mean: ", mean)
 print("population sigma: ", sigma)
 
+# Let the test pass if baseline number is NaN (which happened in
+# the past when we didn't have logic for catching NaN numbers)
+if math.isnan(mean) or math.isnan(sigma):
+    mean = sys.maxsize
+    sigma = 0.001
+
 sample_stats_data = json.loads(args.sample_stats)
 
-sample_mean = sample_stats_data['mean']
-sample_sigma = sample_stats_data['sigma']
+sample_mean = float(sample_stats_data['mean'])
+sample_sigma = float(sample_stats_data['sigma'])
 
 print("sample mean: ", sample_mean)
 print("sample sigma: ", sample_sigma)
 
+if math.isnan(sample_mean):
+    raise Exception('''Error: sample mean is NaN''')
+elif math.isnan(sample_sigma):
+    raise Exception('''Error: sample sigma is NaN''')
+
 z_value = (sample_mean - mean) / sigma
 
 print("z-value: ", z_value)

.jenkins/pytorch/perf_test/test_gpu_speed_mnist.sh

@@ -20,6 +20,9 @@ test_gpu_speed_mnist () {
   SAMPLE_ARRAY=()
   NUM_RUNS=$1
 
+  # Needs warm up to get accurate number
+  python main.py --epochs 1 --no-log
+
   for (( i=1; i<=$NUM_RUNS; i++ )) do
     runtime=$(get_runtime_of_command python main.py --epochs 1 --no-log)
     echo $runtime

.jenkins/pytorch/test.sh

@@ -1,32 +1,57 @@
 #!/bin/bash
 
-COMPACT_JOB_NAME="${BUILD_ENVIRONMENT}-test"
-source "$(dirname "${BASH_SOURCE[0]}")/common.sh"
-
 # Required environment variable: $BUILD_ENVIRONMENT
 # (This is set by default in the Docker images we build, so you don't
 # need to set it yourself.
 
+COMPACT_JOB_NAME="${BUILD_ENVIRONMENT}-test"
+source "$(dirname "${BASH_SOURCE[0]}")/common.sh"
+
 echo "Testing pytorch"
 
-if [[ "$BUILD_ENVIRONMENT" == *rocm* ]]; then
-  echo "Skipping ROCm tests for now"
-  exit 0
+if [ -n "${IN_CIRCLECI}" ]; then
+  if [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda9-* ]]; then
+    # TODO: move this to Docker
+    sudo apt-get -qq update
+    sudo apt-get -qq install --allow-downgrades --allow-change-held-packages libnccl-dev=2.2.13-1+cuda9.0 libnccl2=2.2.13-1+cuda9.0
+  fi
+
+  if [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda8-* ]] || [[ "$BUILD_ENVIRONMENT" == *-xenial-cuda9-cudnn7-py2* ]]; then
+    # TODO: move this to Docker
+    sudo apt-get -qq update
+    sudo apt-get -qq install --allow-downgrades --allow-change-held-packages openmpi-bin libopenmpi-dev
+    sudo apt-get -qq install --no-install-recommends openssh-client openssh-server
+    sudo mkdir -p /var/run/sshd
+  fi
 fi
 
-# JIT C++ extensions require ninja.
-git clone https://github.com/ninja-build/ninja --quiet
-pushd ninja
-python ./configure.py --bootstrap
-export PATH="$PWD:$PATH"
-popd
+# --user breaks ppc64le builds and these packages are already in ppc64le docker
+if [[ "$BUILD_ENVIRONMENT" != *ppc64le* ]]; then
+  # JIT C++ extensions require ninja.
+  pip install -q ninja --user
+  # ninja is installed in /var/lib/jenkins/.local/bin
+  export PATH="/var/lib/jenkins/.local/bin:$PATH"
+
+  # TODO: move this to Docker
+  pip install -q hypothesis --user
+fi
 
 # DANGER WILL ROBINSON.  The LD_PRELOAD here could cause you problems
 # if you're not careful.  Check this if you made some changes and the
 # ASAN test is not working
 if [[ "$BUILD_ENVIRONMENT" == *asan* ]]; then
-    export ASAN_OPTIONS=detect_leaks=0:symbolize=1
-    export UBSAN_OPTIONS=print_stacktrace=1
+    export ASAN_OPTIONS=detect_leaks=0:symbolize=1:strict_init_order=true
+    # We suppress the vptr volation, since we have separate copies of
+    # libprotobuf in both libtorch.so and libcaffe2.so, and it causes
+    # the following problem:
+    #    test_cse (__main__.TestJit) ... torch/csrc/jit/export.cpp:622:38:
+    #        runtime error: member call on address ... which does not point
+    #        to an object of type 'google::protobuf::MessageLite'
+    #        ...: note: object is of type 'onnx_torch::ModelProto'
+    #
+    # This problem should be solved when libtorch.so and libcaffe2.so are
+    # merged.
+    export UBSAN_OPTIONS=print_stacktrace=1:suppressions=$PWD/ubsan.supp
     export PYTORCH_TEST_WITH_ASAN=1
     export PYTORCH_TEST_WITH_UBSAN=1
     # TODO: Figure out how to avoid hard-coding these paths
@@ -49,13 +74,16 @@ if [[ "$BUILD_ENVIRONMENT" == *asan* ]]; then
     (cd test && ! get_exit_code python -c "import torch; torch._C._crash_if_aten_asan(3)")
 fi
 
-export ATEN_DISABLE_AVX=
-export ATEN_DISABLE_AVX2=
-if [[ "${JOB_BASE_NAME}" == *-NO_AVX-* ]]; then
-  export ATEN_DISABLE_AVX=1
+if [[ "$BUILD_ENVIRONMENT" == *rocm* ]]; then
+  export PYTORCH_TEST_WITH_ROCM=1
+  export LANG=C.UTF-8
+  export LC_ALL=C.UTF-8
 fi
-if [[ "${JOB_BASE_NAME}" == *-NO_AVX2-* ]]; then
-  export ATEN_DISABLE_AVX2=1
+
+if [[ "${JOB_BASE_NAME}" == *-NO_AVX-* ]]; then
+  export ATEN_CPU_CAPABILITY=default
+elif [[ "${JOB_BASE_NAME}" == *-NO_AVX2-* ]]; then
+  export ATEN_CPU_CAPABILITY=avx
 fi
 
 test_python_nn() {
@@ -68,14 +96,23 @@ test_python_all_except_nn() {
 
 test_aten() {
   # Test ATen
-  if [[ "$BUILD_ENVIRONMENT" != *asan* ]]; then
+  # The following test(s) of ATen have already been skipped by caffe2 in rocm environment:
+  # scalar_tensor_test, basic, native_test
+  if ([[ "$BUILD_ENVIRONMENT" != *asan* ]] && [[ "$BUILD_ENVIRONMENT" != *rocm* ]]); then
     echo "Running ATen tests with pytorch lib"
     TORCH_LIB_PATH=$(python -c "import site; print(site.getsitepackages()[0])")/torch/lib
     # NB: the ATen test binaries don't have RPATH set, so it's necessary to
     # put the dynamic libraries somewhere were the dynamic linker can find them.
     # This is a bit of a hack.
-    ln -s "$TORCH_LIB_PATH"/libcaffe2* build/bin
-    ln -s "$TORCH_LIB_PATH"/libnccl* build/bin
+    if [[ "$BUILD_ENVIRONMENT" == *ppc64le* ]]; then
+      SUDO=sudo
+    fi
+
+    ${SUDO} ln -s "$TORCH_LIB_PATH"/libc10* build/bin
+    ${SUDO} ln -s "$TORCH_LIB_PATH"/libcaffe2* build/bin
+    ${SUDO} ln -s "$TORCH_LIB_PATH"/libmkldnn* build/bin
+    ${SUDO} ln -s "$TORCH_LIB_PATH"/libnccl* build/bin
+
     ls build/bin
     aten/tools/run_tests.sh build/bin
   fi
@@ -95,7 +132,7 @@ test_torchvision() {
   # this should be a transient requirement...)
   # See https://github.com/pytorch/pytorch/issues/7525
   #time python setup.py install
-  pip install .
+  pip install -q --user .
   popd
 }
 
@@ -104,28 +141,46 @@ test_libtorch() {
      echo "Testing libtorch"
      CPP_BUILD="$PWD/../cpp-build"
      if [[ "$BUILD_ENVIRONMENT" == *cuda* ]]; then
-       "$CPP_BUILD"/libtorch/bin/test_jit
+       "$CPP_BUILD"/caffe2/bin/test_jit
      else
-       "$CPP_BUILD"/libtorch/bin/test_jit "[cpu]"
+       "$CPP_BUILD"/caffe2/bin/test_jit "[cpu]"
      fi
-     python tools/download_mnist.py --quiet -d test/cpp/api/mnist
-     OMP_NUM_THREADS=2 "$CPP_BUILD"/libtorch/bin/test_api
+     python tools/download_mnist.py --quiet -d mnist
+     OMP_NUM_THREADS=2 "$CPP_BUILD"/caffe2/bin/test_api
+  fi
+}
+
+test_custom_script_ops() {
+  if [[ "$BUILD_TEST_LIBTORCH" == "1" ]]; then
+    echo "Testing custom script operators"
+    CUSTOM_OP_BUILD="$PWD/../custom-op-build"
+    pushd test/custom_operator
+    cp -r "$CUSTOM_OP_BUILD" build
+    # Run tests Python-side and export a script module.
+    python test_custom_ops.py -v
+    python model.py --export-script-module=model.pt
+    # Run tests C++-side and load the exported script module.
+    build/test_custom_ops ./model.pt
+    popd
   fi
 }
 
 if [ -z "${JOB_BASE_NAME}" ] || [[ "${JOB_BASE_NAME}" == *-test ]]; then
+  test_torchvision
   test_python_nn
   test_python_all_except_nn
   test_aten
-  test_torchvision
   test_libtorch
+  test_custom_script_ops
 else
   if [[ "${JOB_BASE_NAME}" == *-test1 ]]; then
+    test_torchvision
     test_python_nn
   elif [[ "${JOB_BASE_NAME}" == *-test2 ]]; then
+    test_torchvision
     test_python_all_except_nn
     test_aten
-    test_torchvision
     test_libtorch
+    test_custom_script_ops
   fi
 fi

.jenkins/pytorch/win-build.sh

@@ -38,7 +38,7 @@ EOL
 
 cat >ci_scripts/build_pytorch.bat <<EOL
 
-set PATH=C:\\Program Files\\CMake\\bin;C:\\Program Files\\7-Zip;C:\\curl-7.57.0-win64-mingw\\bin;C:\\Program Files\\Git\\cmd;C:\\Program Files\\Amazon\\AWSCLI;%PATH%
+set PATH=C:\\Program Files\\CMake\\bin;C:\\Program Files\\7-Zip;C:\\ProgramData\\chocolatey\\bin;C:\\Program Files\\Git\\cmd;C:\\Program Files\\Amazon\\AWSCLI;%PATH%
 
 :: Install MKL
 if "%REBUILD%"=="" (
@@ -55,11 +55,11 @@ set LIB=%cd%\\mkl\\lib;%LIB
 :: Install MAGMA
 if "%REBUILD%"=="" (
   if "%BUILD_ENVIRONMENT%"=="" (
-    curl -k https://s3.amazonaws.com/ossci-windows/magma_cuda90_release_mkl_2018.2.185.7z --output magma_cuda90_release_mkl_2018.2.185.7z
+    curl -k https://s3.amazonaws.com/ossci-windows/magma_2.4.0_cuda90_release.7z --output magma_2.4.0_cuda90_release.7z
   ) else (
-    aws s3 cp s3://ossci-windows/magma_cuda90_release_mkl_2018.2.185.7z magma_cuda90_release_mkl_2018.2.185.7z --quiet
+    aws s3 cp s3://ossci-windows/magma_2.4.0_cuda90_release.7z magma_2.4.0_cuda90_release.7z --quiet
   )
-  7z x -aoa magma_cuda90_release_mkl_2018.2.185.7z -omagma
+  7z x -aoa magma_2.4.0_cuda90_release.7z -omagma
 )
 set MAGMA_HOME=%cd%\\magma
 
@@ -80,18 +80,29 @@ if "%REBUILD%"=="" (
 )
 
 :: Install Miniconda3
-if "%REBUILD%"=="" (
-  IF EXIST C:\\Jenkins\\Miniconda3 ( rd /s /q C:\\Jenkins\\Miniconda3 )
-  curl -k https://repo.continuum.io/miniconda/Miniconda3-latest-Windows-x86_64.exe -O
-  .\Miniconda3-latest-Windows-x86_64.exe /InstallationType=JustMe /RegisterPython=0 /S /AddToPath=0 /D=C:\\Jenkins\\Miniconda3
+if "%BUILD_ENVIRONMENT%"=="" (
+  set CONDA_PARENT_DIR=%CD%
+) else (
+  set CONDA_PARENT_DIR=C:\\Jenkins
+)
+if "%REBUILD%"=="" (
+  IF EXIST %CONDA_PARENT_DIR%\\Miniconda3 ( rd /s /q %CONDA_PARENT_DIR%\\Miniconda3 )
+  curl -k https://repo.continuum.io/miniconda/Miniconda3-latest-Windows-x86_64.exe -O
+  .\Miniconda3-latest-Windows-x86_64.exe /InstallationType=JustMe /RegisterPython=0 /S /AddToPath=0 /D=%CONDA_PARENT_DIR%\\Miniconda3
+)
+call %CONDA_PARENT_DIR%\\Miniconda3\\Scripts\\activate.bat %CONDA_PARENT_DIR%\\Miniconda3
+if "%REBUILD%"=="" (
+  :: We have to pin Python version to 3.6.7, until mkl supports Python 3.7
+  call conda install -y -q python=3.6.7 numpy cffi pyyaml boto3
 )
-call C:\\Jenkins\\Miniconda3\\Scripts\\activate.bat C:\\Jenkins\\Miniconda3
-if "%REBUILD%"=="" ( call conda install -y -q numpy cffi pyyaml boto3 )
 
 :: Install ninja
 if "%REBUILD%"=="" ( pip install ninja )
 
+set WORKING_DIR=%CD%
+call "C:\\Program Files (x86)\\Microsoft Visual Studio\\2017\\Community\\VC\\Auxiliary\\Build\\vcvarsall.bat" x64
 call "C:\\Program Files (x86)\\Microsoft Visual Studio\\2017\\Community\\VC\\Auxiliary\\Build\\vcvarsall.bat" x86_amd64
+cd %WORKING_DIR%
 
 git submodule update --init --recursive
 
@@ -129,7 +140,7 @@ if not "%USE_CUDA%"=="0" (
   if "%REBUILD%"=="" (
     sccache --show-stats
     sccache --zero-stats
-    rd /s /q C:\\Jenkins\\Miniconda3\\Lib\\site-packages\\torch
+    rd /s /q %CONDA_PARENT_DIR%\\Miniconda3\\Lib\\site-packages\\torch
     copy %CD%\\tmp_bin\\sccache.exe tmp_bin\\nvcc.exe
   )
 
@@ -139,9 +150,10 @@ if not "%USE_CUDA%"=="0" (
 
   python setup.py install && sccache --show-stats && (
     if "%BUILD_ENVIRONMENT%"=="" (
-      echo "NOTE: To run \`import torch\`, please make sure to activate the conda environment by running \`call C:\\Jenkins\\Miniconda3\\Scripts\\activate.bat C:\\Jenkins\\Miniconda3\` in Command Prompt before running Git Bash."
+      echo NOTE: To run \`import torch\`, please make sure to activate the conda environment by running \`call %CONDA_PARENT_DIR%\\Miniconda3\\Scripts\\activate.bat %CONDA_PARENT_DIR%\\Miniconda3\` in Command Prompt before running Git Bash.
     ) else (
-      7z a %IMAGE_COMMIT_TAG%.7z C:\\Jenkins\\Miniconda3\\Lib\\site-packages\\torch && python ci_scripts\\upload_image.py %IMAGE_COMMIT_TAG%.7z
+      mv %CD%\\build\\bin\\test_api.exe %CONDA_PARENT_DIR%\\Miniconda3\\Lib\\site-packages\\torch\\lib
+      7z a %IMAGE_COMMIT_TAG%.7z %CONDA_PARENT_DIR%\\Miniconda3\\Lib\\site-packages\\torch && python ci_scripts\\upload_image.py %IMAGE_COMMIT_TAG%.7z
     )
   )
 )

.jenkins/pytorch/win-test.sh

@@ -36,18 +36,29 @@ EOL
 
 cat >ci_scripts/setup_pytorch_env.bat <<EOL
 
-set PATH=C:\\Program Files\\CMake\\bin;C:\\Program Files\\7-Zip;C:\\curl-7.57.0-win64-mingw\\bin;C:\\Program Files\\Git\\cmd;C:\\Program Files\\Amazon\\AWSCLI;%PATH%
+set PATH=C:\\Program Files\\CMake\\bin;C:\\Program Files\\7-Zip;C:\\ProgramData\\chocolatey\\bin;C:\\Program Files\\Git\\cmd;C:\\Program Files\\Amazon\\AWSCLI;%PATH%
 
 :: Install Miniconda3
-IF EXIST C:\\Jenkins\\Miniconda3 ( rd /s /q C:\\Jenkins\\Miniconda3 )
-curl https://repo.continuum.io/miniconda/Miniconda3-latest-Windows-x86_64.exe -O
-.\Miniconda3-latest-Windows-x86_64.exe /InstallationType=JustMe /RegisterPython=0 /S /AddToPath=0 /D=C:\\Jenkins\\Miniconda3
-call C:\\Jenkins\\Miniconda3\\Scripts\\activate.bat C:\\Jenkins\\Miniconda3
-call conda install -y -q numpy mkl cffi pyyaml boto3
-
-pip install ninja
+if "%BUILD_ENVIRONMENT%"=="" (
+    set CONDA_PARENT_DIR=%CD%
+) else (
+    set CONDA_PARENT_DIR=C:\\Jenkins
+)
+if NOT "%BUILD_ENVIRONMENT%"=="" (
+    IF EXIST %CONDA_PARENT_DIR%\\Miniconda3 ( rd /s /q %CONDA_PARENT_DIR%\\Miniconda3 )
+    curl https://repo.continuum.io/miniconda/Miniconda3-latest-Windows-x86_64.exe -O
+    .\Miniconda3-latest-Windows-x86_64.exe /InstallationType=JustMe /RegisterPython=0 /S /AddToPath=0 /D=%CONDA_PARENT_DIR%\\Miniconda3
+)
+call %CONDA_PARENT_DIR%\\Miniconda3\\Scripts\\activate.bat %CONDA_PARENT_DIR%\\Miniconda3
+if NOT "%BUILD_ENVIRONMENT%"=="" (
+    :: We have to pin Python version to 3.6.7, until mkl supports Python 3.7
+    call conda install -y -q python=3.6.7 numpy mkl cffi pyyaml boto3
+)
+pip install ninja future hypothesis
 
+set WORKING_DIR=%CD%
 call "C:\\Program Files (x86)\\Microsoft Visual Studio\\2017\\Community\\VC\\Auxiliary\\Build\\vcvarsall.bat" x86_amd64
+cd %WORKING_DIR%
 
 set PATH=C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v9.0\\bin;C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v9.0\\libnvvp;%PATH%
 set CUDA_PATH=C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v9.0
@@ -58,13 +69,14 @@ set CUDA_TOOLKIT_ROOT_DIR=C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\
 set CUDNN_ROOT_DIR=C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v9.0
 set PYTHONPATH=%CD%\\test;%PYTHONPATH%
 
-cd test/
-
-python ..\\ci_scripts\\download_image.py %IMAGE_COMMIT_TAG%.7z
-
-7z x %IMAGE_COMMIT_TAG%.7z
-
-cd ..
+if NOT "%BUILD_ENVIRONMENT%"=="" (
+    cd test/
+    python ..\\ci_scripts\\download_image.py %IMAGE_COMMIT_TAG%.7z
+    7z x %IMAGE_COMMIT_TAG%.7z
+    cd ..
+) else (
+    xcopy /s %CONDA_PARENT_DIR%\\Miniconda3\\Lib\\site-packages\\torch .\\test\\torch\\
+)
 
 EOL
 
@@ -78,14 +90,47 @@ call ci_scripts/setup_pytorch_env.bat
 cd test/ && python run_test.py --exclude nn --verbose && cd ..
 EOL
 
+cat >ci_scripts/test_custom_script_ops.bat <<EOL
+call ci_scripts/setup_pytorch_env.bat
+
+cd test/custom_operator
+
+:: Build the custom operator library.
+mkdir build
+cd build
+:: Note: Caffe2 does not support MSVC + CUDA + Debug mode (has to be Release mode)
+cmake -DCMAKE_PREFIX_PATH=%CD%\\..\\..\\torch -DCMAKE_BUILD_TYPE=Release -GNinja ..
+ninja -v
+cd ..
+
+:: Run tests Python-side and export a script module.
+python test_custom_ops.py -v
+python model.py --export-script-module="build/model.pt"
+:: Run tests C++-side and load the exported script module.
+cd build
+set PATH=C:\\Program Files\\NVIDIA Corporation\\NvToolsExt/bin/x64;%CD%\\..\\..\\torch\\lib;%PATH%
+test_custom_ops.exe model.pt
+EOL
+
+cat >ci_scripts/test_libtorch.bat <<EOL
+call ci_scripts/setup_pytorch_env.bat
+dir
+dir %CD%\\test 
+dir %CD%\\test\\torch
+dir %CD%\\test\\torch\\lib
+cd %CD%\\test\\torch\\lib
+set PATH=C:\\Program Files\\NVIDIA Corporation\\NvToolsExt/bin/x64;%CD%\\..\\..\\torch\\lib;%PATH%
+test_api.exe --gtest_filter="-IntegrationTest.MNIST*"
+EOL
+
 run_tests() {
     if [ -z "${JOB_BASE_NAME}" ] || [[ "${JOB_BASE_NAME}" == *-test ]]; then
-        ci_scripts/test_python_nn.bat && ci_scripts/test_python_all_except_nn.bat
+        ci_scripts/test_python_nn.bat && ci_scripts/test_python_all_except_nn.bat && ci_scripts/test_custom_script_ops.bat && ci_scripts/test_libtorch.bat
     else
         if [[ "${JOB_BASE_NAME}" == *-test1 ]]; then
             ci_scripts/test_python_nn.bat
         elif [[ "${JOB_BASE_NAME}" == *-test2 ]]; then
-            ci_scripts/test_python_all_except_nn.bat
+            ci_scripts/test_python_all_except_nn.bat && ci_scripts/test_custom_script_ops.bat && ci_scripts/test_libtorch.bat
         fi
     fi
 }

.travis.yml

@@ -16,7 +16,28 @@ matrix:
         python: "2.7"
         install: pip install flake8
         script: flake8
+      - env: LINT_CHECK
+        python: "3.7"
+        dist: xenial    # required for Python 3.7 (travis-ci/travis-ci#9069)
+        sudo: required  # required for Python 3.7 (travis-ci/travis-ci#9069)
+        install: pip install flake8
+        script: flake8
       - env: MYPY_TYPE_CHECK
         python: "3.6"
         install: pip install mypy mypy-extensions
         script: mypy @mypy-files.txt
+      - env: CPP_DOC_CHECK
+        python: "3.6"
+        install: 
+          - sudo apt-get install -y doxygen
+          - pip install -r requirements.txt
+        script: cd docs/cpp/source && ./check-doxygen.sh
+      - env: CLANG_TIDY
+        python: "3.6"
+        addons:
+          apt:
+            sources:
+              - ubuntu-toolchain-r-test
+              - llvm-toolchain-trusty
+            packages: clang-tidy
+        script: tools/run-clang-tidy-in-ci.sh

CMakeLists.txt

@@ -5,11 +5,14 @@ cmake_minimum_required(VERSION 3.5 FATAL_ERROR)
 # ---[ Project and semantic versioning.
 project(Caffe2 CXX C)
 
-set(CAFFE2_VERSION_MAJOR 0)
-set(CAFFE2_VERSION_MINOR 8)
-set(CAFFE2_VERSION_PATCH 2)
-set(CAFFE2_VERSION
-    "${CAFFE2_VERSION_MAJOR}.${CAFFE2_VERSION_MINOR}.${CAFFE2_VERSION_PATCH}")
+set(CMAKE_INSTALL_MESSAGE NEVER)
+
+set(CMAKE_CXX_STANDARD 11)
+if (NOT MSVC)
+  set(CMAKE_C_STANDARD 11)
+endif()
+
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
 # One variable that determines whether the current cmake process is being run
 # with the main Caffe2 library. This is useful for building modules - if
@@ -53,12 +56,16 @@ endif()
 # Note to developers: if you add an option below, make sure you also add it to
 # cmake/Summary.cmake so that the summary prints out the option values.
 include(CMakeDependentOption)
-option(BUILD_CAFFE2 "Build Caffe2" ON)
-option(BUILD_ATEN "Build ATen" OFF)
-option(BUILD_BINARY "Build C++ binaries" ON)
+option(BUILD_TORCH "Build Torch" OFF)
+option(ATEN_NO_TEST "Do not build ATen test binaries" OFF)
+option(BUILD_ATEN_MOBILE "Build ATen for Android and iOS" OFF)
+option(BUILD_ATEN_ONLY "Build only a subset focused on ATen only" OFF)
+option(BUILD_BINARY "Build C++ binaries" OFF)
 option(BUILD_DOCS "Build Caffe2 documentation" OFF)
 option(BUILD_CUSTOM_PROTOBUF "Build and use Caffe2's own protobuf under third_party" ON)
 option(BUILD_PYTHON "Build Python binaries" ON)
+option(BUILD_CAFFE2_OPS "Build Caffe2 operators" ON)
+option(BUILD_C10_EXPERIMENTAL_OPS "Build c10 experimental operators" ON)
 option(BUILD_SHARED_LIBS "Build libcaffe2.so" ON)
 cmake_dependent_option(
     CAFFE2_LINK_LOCAL_PROTOBUF "If set, build protobuf inside libcaffe2.so." ON
@@ -66,95 +73,107 @@ cmake_dependent_option(
 cmake_dependent_option(
     CAFFE2_USE_MSVC_STATIC_RUNTIME "Using MSVC static runtime libraries" ON
     "NOT BUILD_SHARED_LIBS" OFF)
-cmake_dependent_option(
-    BUILD_TEST "Build Caffe2 C++ test binaries (need gtest and gbenchmark)" OFF
-    "BUILD_CAFFE2" OFF)
+option(BUILD_TEST "Build C++ test binaries (need gtest and gbenchmark)" OFF)
 cmake_dependent_option(
     INSTALL_TEST "Install test binaries if BUILD_TEST is on" OFF
     "BUILD_TEST" OFF)
 option(USE_ACL "Use ARM Compute Library" OFF)
 option(USE_ASAN "Use Address Sanitizer" OFF)
-option(USE_ATEN "Use ATen" OFF)
 option(USE_CUDA "Use CUDA" ON)
-option(USE_ROCM "Use ROCm" OFF)
+option(USE_ROCM "Use ROCm" ON)
 option(CAFFE2_STATIC_LINK_CUDA "Statically link CUDA libraries" OFF)
 cmake_dependent_option(
     USE_CUDNN "Use cuDNN" ON
     "USE_CUDA" OFF)
+option(USE_FBGEMM "Use FBGEMM (quantized 8-bit server operators)" OFF)
 option(USE_FFMPEG "Use ffmpeg" OFF)
-cmake_dependent_option(
-    USE_GFLAGS "Use GFLAGS" ON
-    "BUILD_CAFFE2" OFF)
-cmake_dependent_option(
-    USE_GLOG "Use GLOG" ON
-    "BUILD_CAFFE2" OFF)
-cmake_dependent_option(
-    USE_GLOO "Use Gloo" ON
-    "BUILD_CAFFE2" OFF)
-option(USE_GLOO_IBVERBS "Use Gloo IB verbs for distributed support" OFF)
-cmake_dependent_option(
-    USE_LEVELDB "Use LEVELDB" ON
-    "BUILD_CAFFE2" OFF)
+option(USE_GFLAGS "Use GFLAGS" ON)
+option(USE_GLOG "Use GLOG" ON)
+option(USE_LEVELDB "Use LEVELDB" ON)
 option(USE_LITE_PROTO "Use lite protobuf instead of full." OFF)
-cmake_dependent_option(
-    USE_LMDB "Use LMDB" ON
-    "BUILD_CAFFE2" OFF)
-cmake_dependent_option(
-    USE_METAL "Use Metal for iOS build" ON
-    "BUILD_CAFFE2" OFF)
-cmake_dependent_option(
-    USE_MOBILE_OPENGL "Use OpenGL for mobile code" ON
-    "BUILD_CAFFE2" OFF)
-cmake_dependent_option(
-    USE_MPI "Use MPI" ON
-    "BUILD_CAFFE2" OFF)
+option(USE_LMDB "Use LMDB" ON)
+option(USE_METAL "Use Metal for iOS build" ON)
+option(USE_MOBILE_OPENGL "Use OpenGL for mobile code" ON)
 option(USE_NATIVE_ARCH "Use -march=native" OFF)
 option(USE_NCCL "Use NCCL" ON)
 option(USE_SYSTEM_NCCL "Use system-wide NCCL" OFF)
-option(USE_NERVANA_GPU "Use Nervana GPU backend" OFF)
 option(USE_NNAPI "Use NNAPI" OFF)
 option(USE_NNPACK "Use NNPACK" ON)
 option(USE_NUMA "Use NUMA (only available on Linux)" ON)
 cmake_dependent_option(
     USE_NVRTC "Use NVRTC. Only available if USE_CUDA is on." OFF
     "USE_CUDA" OFF)
+option(USE_NUMPY "Use NumPy" ON)
 option(USE_OBSERVERS "Use observers module." OFF)
 option(USE_OPENCL "Use OpenCL" OFF)
-cmake_dependent_option(
-    USE_OPENCV "Use OpenCV" ON
-    "BUILD_CAFFE2" OFF)
+option(USE_OPENCV "Use OpenCV" ON)
 option(USE_OPENMP "Use OpenMP for parallel code" OFF)
 option(USE_PROF "Use profiling" OFF)
+option(USE_QNNPACK "Use QNNPACK (quantized 8-bit operators)" ON)
 option(USE_REDIS "Use Redis" OFF)
 option(USE_ROCKSDB "Use RocksDB" OFF)
 option(USE_SNPE "Use Qualcomm's SNPE library" OFF)
+option(USE_SYSTEM_EIGEN_INSTALL
+    "Use system Eigen instead of the one under third_party" OFF)
 option(USE_TENSORRT "Using Nvidia TensorRT library" OFF)
 option(USE_ZMQ "Use ZMQ" OFF)
 option(USE_ZSTD "Use ZSTD" OFF)
 option(USE_MKLDNN "Use MKLDNN" OFF)
+option(USE_DISTRIBUTED "Use distributed" ON)
 cmake_dependent_option(
-  USE_IDEEP "Use IDEEP interface in MKL BLAS" ON
-    "BUILD_CAFFE2" OFF)
+    USE_MPI "Use MPI for Caffe2. Only available if USE_DISTRIBUTED is on." ON
+    "USE_DISTRIBUTED" OFF)
 cmake_dependent_option(
-  USE_MKLML "Use MKLML interface in MKL BLAS" ON
-    "BUILD_CAFFE2" OFF)
-option(USE_DISTRIBUTED "Use THD (distributed)" OFF)
-option(USE_DISTRIBUTED_MW "Use THD (distributed) master worker" OFF)
+    USE_GLOO "Use Gloo. Only available if USE_DISTRIBUTED is on." ON
+    "USE_DISTRIBUTED" OFF)
+cmake_dependent_option(
+    USE_GLOO_IBVERBS "Use Gloo IB verbs for distributed. Only available if USE_GLOO is on." OFF
+    "USE_GLOO" OFF)
 
 # Used when building Caffe2 through setup.py
 option(BUILDING_WITH_TORCH_LIBS "Tell cmake if Caffe2 is being built alongside torch libs" OFF)
 
-if (USE_ATEN)
-    set(BUILD_ATEN ${USE_ATEN})
+SET(ONNX_NAMESPACE "onnx_c2" CACHE STRING "onnx namespace")
+
+if (ANDROID OR IOS)
+  set(BUILD_ATEN_MOBILE ON)
 endif()
 
+if (BUILD_ATEN_ONLY)
+  set(BUILD_CAFFE2_OPS OFF)
+  set(BUILD_PYTHON OFF)
+  set(USE_NUMA OFF)
+  set(USE_LEVELDB OFF)
+  set(USE_GFLAGS OFF)
+  set(USE_GLOG OFF)
+  set(USE_NCCL OFF)
+  set(USE_NNPACK OFF)
+  set(USE_NUMPY OFF)
+  set(USE_OPENCV OFF)
+  set(USE_MKLDNN OFF)
+  set(USE_DISTRIBUTED OFF)
+  set(USE_LMDB OFF)
+endif()
+
+# ---[ Utils
+# TODO: merge the following 3 files into cmake/public/utils.cmake.
+include(cmake/Utils.cmake)
+include(cmake/public/utils.cmake)
+
+# ---[ Version numbers for generated libraries
+set(TORCH_DEFAULT_VERSION "1.0.0")
+set(TORCH_BUILD_VERSION "${TORCH_DEFAULT_VERSION}" CACHE STRING "Torch build version")
+if (NOT TORCH_BUILD_VERSION)
+  # An empty string was specified so force version to the default
+  set(TORCH_BUILD_VERSION "${TORCH_DEFAULT_VERSION}"
+    CACHE STRING "Torch build version" FORCE)
+endif()
+caffe2_parse_version_str(TORCH ${TORCH_BUILD_VERSION})
+caffe2_parse_version_str(CAFFE2 ${TORCH_BUILD_VERSION})
+
 # ---[ CMake scripts + modules
 list(APPEND CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/cmake/Modules)
 
-if (MSVC AND ${BUILD_SHARED_LIBS})
-  set(CMAKE_WINDOWS_EXPORT_ALL_SYMBOLS ON)
-endif()
-
 # ---[ CMake build directories
 set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib)
 set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib)
@@ -178,11 +197,6 @@ include(cmake/MiscCheck.cmake)
 # External projects
 include(ExternalProject)
 
-# ---[ Utils
-# TODO: merge the following 3 files into cmake/public/utils.cmake.
-include(cmake/Utils.cmake)
-include(cmake/public/utils.cmake)
-
 # ---[ Dependencies
 include(cmake/Dependencies.cmake)
 
@@ -214,11 +228,18 @@ if(NOT MSVC)
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-variable")
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-function")
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-result")
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-strict-overflow")
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-strict-aliasing")
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-error=deprecated-declarations")
+  if (CMAKE_COMPILER_IS_GNUCXX AND NOT (CMAKE_CXX_COMPILER_VERSION VERSION_LESS 7.0.0))
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-stringop-overflow")
+  endif()
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-error=pedantic")
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-error=redundant-decls")
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-error=old-style-cast")
   # These flags are not available in GCC-4.8.5. Set only when using clang.
   # Compared against https://gcc.gnu.org/onlinedocs/gcc-4.8.5/gcc/Option-Summary.html
-  if ("${CMAKE_CXX_COMPILER_ID}" STREQUAL "Clang")
+  if ("${CMAKE_CXX_COMPILER_ID}" MATCHES "Clang")
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-invalid-partial-specialization")
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-typedef-redefinition")
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unknown-warning-option")
@@ -228,15 +249,20 @@ if(NOT MSVC)
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-c++14-extensions")
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-constexpr-not-const")
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-missing-braces")
+    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()
   if ($ENV{WERROR})
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Werror")
   endif($ENV{WERROR})
+  if (NOT APPLE)
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-but-set-variable")
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-maybe-uninitialized")
+  endif()
 else()
   foreach(flag_var
       CMAKE_CXX_FLAGS CMAKE_CXX_FLAGS_DEBUG CMAKE_CXX_FLAGS_RELEASE
@@ -263,6 +289,21 @@ if (USE_ASAN)
     set (CMAKE_LINKER_FLAGS_DEBUG "${CMAKE_STATIC_LINKER_FLAGS_DEBUG} -fsanitize=address")
 endif()
 
+if (APPLE)
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-private-field")
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-missing-braces")
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-c++14-extensions")
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-constexpr-not-const")
+endif()
+
+if (EMSCRIPTEN)
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-implicit-function-declaration -DEMSCRIPTEN -s DISABLE_EXCEPTION_CATCHING=0")
+endif()
+
+if(CMAKE_COMPILER_IS_GNUCXX AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER 7.0.0)
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-stringop-overflow")
+endif()
+
 if(ANDROID)
   if(CMAKE_COMPILER_IS_GNUCXX)
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -s")
@@ -286,12 +327,10 @@ include_directories(BEFORE ${PROJECT_SOURCE_DIR})
 # in PROJECT_SOURCE_DIR.
 include_directories(BEFORE ${PROJECT_BINARY_DIR})
 
-# ---[ Old caffe protobuf
-if(BUILD_CAFFE2)
-  add_subdirectory(caffe/proto)
-endif()
+include_directories(BEFORE ${PROJECT_SOURCE_DIR}/aten/src/)
 
 # ---[ Main build
+add_subdirectory(c10)
 add_subdirectory(caffe2)
 
 # --[ Documentation
@@ -308,7 +347,7 @@ if(BUILD_DOCS)
 
     if(EXISTS ${CMAKE_CURRENT_BINARY_DIR}/docs)
       file(REMOVE_RECURSE ${CMAKE_CURRENT_BINARY_DIR}/docs)
-    endif (EXISTS ${CMAKE_CURRENT_BINARY_DIR}/docs)
+    endif()
 
     file(MAKE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/docs)
     configure_file(${DOXYGEN_C_IN} ${DOXYGEN_C_OUT} @ONLY)
@@ -325,10 +364,10 @@ if(BUILD_DOCS)
         WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
         COMMENT "Generating Python API documentation with Doxygen"
         VERBATIM)
-  else (DOXYGEN_FOUND)
+  else()
     message(FATAL_ERROR "Doxygen needs to be installed to generate the documentation")
-  endif (DOXYGEN_FOUND)
-endif (BUILD_DOCS)
+  endif()
+endif()
 
 # ---[ CMake related files
 # Uninistall option.
@@ -383,6 +422,8 @@ if (BUILD_SHARED_LIBS)
       ${PROJECT_SOURCE_DIR}/cmake/public/cuda.cmake
       ${PROJECT_SOURCE_DIR}/cmake/public/glog.cmake
       ${PROJECT_SOURCE_DIR}/cmake/public/gflags.cmake
+      ${PROJECT_SOURCE_DIR}/cmake/public/mkl.cmake
+      ${PROJECT_SOURCE_DIR}/cmake/public/mkldnn.cmake
       ${PROJECT_SOURCE_DIR}/cmake/public/protobuf.cmake
       ${PROJECT_SOURCE_DIR}/cmake/public/threads.cmake
       ${PROJECT_SOURCE_DIR}/cmake/public/utils.cmake
@@ -402,19 +443,15 @@ else()
 endif()
 
 # ---[ Modules
-if (BUILD_CAFFE2)
-  add_subdirectory(modules)
-endif()
+add_subdirectory(modules)
 
 # ---[ Binaries
 # Binaries will be built after the Caffe2 main libraries and the modules
 # are built. For the binaries, they will be linked to the Caffe2 main
 # libraries, as well as all the modules that are built with Caffe2 (the ones
 # built in the previous Modules section above).
-if (BUILD_CAFFE2)
-  if (BUILD_BINARY)
-    add_subdirectory(binaries)
-  endif()
+if (BUILD_BINARY)
+  add_subdirectory(binaries)
 endif()
 
 include(cmake/Summary.cmake)

CODEOWNERS

@@ -1,21 +1,9 @@
 # This is a comment.
 # Each line is a file pattern followed by one or more owners.
 
-/aten/ @apaszke @soumith @colesbury @gchanan @zdevito @ezyang
-/torch/ @apaszke @soumith @colesbury @gchanan @zdevito @ezyang
-/docs/source @apaszke @soumith @colesbury @gchanan @zdevito @ezyang @ssnl @zou3519
-/test @apaszke @soumith @colesbury @gchanan @zdevito @ezyang
-/tools @apaszke @soumith @colesbury @gchanan @zdevito @ezyang
-/README.md @apaszke @soumith @colesbury @gchanan @zdevito @ezyang
-/setup.py @apaszke @soumith @colesbury @gchanan @zdevito @ezyang
-/requirements.txt @apaszke @soumith @colesbury @gchanan @zdevito @ezyang
-/torch/csrc/api/ @apaszke @soumith @colesbury @gchanan @zdevito @ezyang @ebetica @goldsborough
-/test/cpp/api/ @apaszke @soumith @colesbury @gchanan @zdevito @ezyang @ebetica @goldsborough
-/torch/onnx/ @anderspapitto @bddppq @dzhulgakov @ezyang @houseroad @jamesr66a @smessmer @Yangqing
-/torch/csrc/onnx/ @anderspapitto @bddppq @dzhulgakov @ezyang @houseroad @jamesr66a @smessmer @Yangqing
-/torch/csrc/jit/passes/onnx/ @anderspapitto @bddppq @dzhulgakov @ezyang @houseroad @jamesr66a @smessmer @Yangqing
-/test/onnx/ @anderspapitto @bddppq @dzhulgakov @ezyang @houseroad @jamesr66a @smessmer @Yangqing
-/scripts/onnx/ @anderspapitto @bddppq @dzhulgakov @ezyang @houseroad @jamesr66a @smessmer @Yangqing
+/docs/cpp @goldsborough @ebetica
+/torch/csrc/api/ @ebetica @goldsborough
+/test/cpp/api/ @ebetica @goldsborough
 /torch/lib/c10d/ @apaszke @pietern @teng-li
 /torch/csrc/distributed/ @apaszke @pietern @teng-li
 /torch/distributed/ @apaszke @pietern @teng-li

CONTRIBUTING.md

@@ -19,18 +19,18 @@ If you are not familiar with creating a Pull Request, here are some guides:
 - https://help.github.com/articles/creating-a-pull-request/
 
 
-## Developing locally with PyTorch
+## Developing PyTorch
 
-To locally develop with PyTorch, here are some tips:
+To develop PyTorch on your machine, here are some tips:
 
-1. Uninstall all existing pytorch installs
+1. Uninstall all existing PyTorch installs:
 ```
 conda uninstall pytorch
 pip uninstall torch
 pip uninstall torch # run this command twice
 ```
 
-2. Locally clone a copy of PyTorch from source:
+2. Clone a copy of PyTorch from source:
 
 ```
 git clone https://github.com/pytorch/pytorch
@@ -75,6 +75,84 @@ You do not need to repeatedly install after modifying python files.
 In case you want to reinstall, make sure that you uninstall pytorch first by running `pip uninstall torch`
 and `python setup.py clean`. Then you can install in `build develop` mode again.
 
+## Codebase structure
+
+* [c10](c10) - Core library files that work everywhere, both server
+  and mobile.  We are slowly moving pieces from ATen/core here.
+  This library is intended only to contain essential functionality,
+  and appropriate to use in settings where binary size matters.  (But
+  you'll have a lot of missing functionality if you try to use it
+  directly.)
+* [aten](aten) - C++ tensor library for PyTorch (no autograd support)
+  * src
+    * [TH](aten/src/TH)
+      [THC](aten/src/THC)
+      [THNN](aten/src/THNN)
+      [THCUNN](aten/src/THCUNN) - Legacy library code from the original
+      Torch.  Try not to add things here; we're slowly porting these to
+      native.
+      * generic - Contains actual implementations of operators,
+        parametrized over `scalar_t`.  Files here get compiled N times
+        per supported scalar type in PyTorch.
+    * ATen
+      * [core](aten/src/ATen/core) - Core functionality of ATen.  This
+        is migrating to top-level c10 folder.
+      * [native](aten/src/ATen/native) - Modern implementations of
+        operators.  If you want to write a new operator, here is where
+        it should go.  Most CPU operators go in the top level directory,
+        except for operators which need to be compiled specially; see
+        cpu below.
+        * [cpu](aten/src/ATen/native/cpu) - Not actually CPU
+          implementations of operators, but specifically implementations
+          which are compiled with processor-specific instructions, like
+          AVX.  See the README for more details.
+        * [cuda](aten/src/ATen/native/cuda) - CUDA implementations of
+          operators.
+        * [sparse](aten/src/ATen/native/sparse) - CPU and CUDA
+          implementations of COO sparse tensor operations
+        * [mkl](aten/src/ATen/native/mkl) [mkldnn](aten/src/ATen/native/mkldnn)
+          [miopen](aten/src/ATen/native/miopen) [cudnn](aten/src/ATen/native/cudnn)
+          - implementations of operators which simply bind to some
+            backend library.
+* [torch](torch) - The actual PyTorch library.  Everything that is not
+  in csrc is Python modules, following the PyTorch Python frontend
+  module structure.
+  * [csrc](torch/csrc) - C++ files composing the PyTorch library.  Files
+    in this directory tree are a mix of Python binding code, and C++
+    heavy lifting.  Consult `setup.py` for the canonical list of Python
+    binding files; conventionally, they are often prefixed with
+    `python_`.
+    * [jit](torch/csrc/jit) - Compiler and frontend for TorchScript JIT
+      frontend.
+    * [autograd](torch/csrc/autograd) - Implementation of reverse-mode automatic
+      differentation
+    * [api](torch/csrc/api) - The PyTorch C++ frontend.
+    * [distributed](torch/csrc/distributed) - Distributed training
+      support for PyTorch.
+* [tools](tools) - Code generation scripts for the PyTorch library.
+  See README of this directory for more details.
+* [test](tests) - Python unit tests for PyTorch Python frontend
+  * [test_torch.py](test/test_torch.py) - Basic tests for PyTorch
+    functionality
+  * [test_autograd.py](test/test_autograd.py) - Tests for non-NN
+    automatic differentiation support
+  * [test_nn.py](test/test_nn.py) - Tests for NN operators and
+    their automatic differentiation
+  * [test_jit.py](test/test_jit.py) - Tests for the JIT compiler
+    and TorchScript
+  * ...
+  * [cpp](test/cpp) - C++ unit tests for PyTorch C++ frontend
+  * [expect](test/expect) - Automatically generated "expect" files
+    which are used to compare against expected output.
+  * [onnx](test/onnx) - Tests for ONNX export functionality,
+    using both PyTorch and Caffe2.
+* [caffe2](caffe2) - The Caffe2 library.
+  * [core](caffe2/core) - Core files of Caffe2, e.g., tensor, workspace,
+    blobs, etc.
+  * [operators](caffe2/operators) - Operators of Caffe2
+  * [python](caffe2/python) - Python bindings to Caffe2
+  * ...
+
 ## Unit testing
 
 PyTorch's testing is located under `test/`. Run the entire test suite with
@@ -104,6 +182,18 @@ PyTorch uses [Google style](http://sphinxcontrib-napoleon.readthedocs.io/en/late
 for formatting docstrings. Length of line inside docstrings block must be limited to 80 characters to
 fit into Jupyter documentation popups.
 
+For C++ documentation (https://pytorch.org/cppdocs), we use
+[Doxygen](http://www.doxygen.nl/) and then convert it to
+[Sphinx](http://www.sphinx-doc.org/) via
+[Breathe](https://github.com/michaeljones/breathe) and
+[Exhale](https://github.com/svenevs/exhale). Check the [Doxygen
+reference](http://www.stack.nl/~dimitri/doxygen/manual/index.html) for more
+information on the documentation syntax. To build the documentation locally,
+`cd` into `docs/cpp` and then `make html`.
+
+We run Doxygen in CI (Travis) to verify that you do not use invalid Doxygen
+commands. To run this check locally, run `./check-doxygen.sh` from inside
+`docs/cpp`.
 
 ## Managing multiple build trees
 
@@ -139,19 +229,20 @@ not very optimized for incremental rebuilds, this will actually be very slow.
 Far better is to only request rebuilds of the parts of the project you are
 working on:
 
-- Working on `torch/csrc`?  Run `python setup.py develop` to rebuild
+- Working on the Python bindings?  Run `python setup.py develop` to rebuild
   (NB: no `build` here!)
 
-- Working on `torch/lib/TH`, did not make any cmake changes, and just want to
-  see if it compiles?  Run `(cd torch/lib/build/TH && make install -j$(getconf _NPROCESSORS_ONLN))`.  This
-  applies for any other subdirectory of `torch/lib`.  **Warning: Changes you
-  make here will not be visible from Python.**  See below.
+- Working on `torch/csrc` or `aten`?  Run `python setup.py rebuild_libtorch` to
+  rebuild and avoid having to rebuild other dependent libraries we
+  depend on.
 
-- Working on `torch/lib` and want to run your changes / rerun cmake?  Run
-  `python setup.py build_deps`.  Note that this will rerun cmake for
-  every subdirectory in TH; if you are only working on one project,
-  consider editing `torch/lib/build_all.sh` and commenting out the
-  `build` lines of libraries you are not working on.
+- Working on one of the other dependent libraries? The other valid
+  targets are listed in `dep_libs` in `setup.py`. prepend `build_` to
+  get a target, and run as e.g. `python setup.py build_gloo`.
+
+- Working on a test binary?  Run `(cd build && ninja bin/test_binary_name)` to
+  rebuild only that test binary (without rerunning cmake).  (Replace `ninja` with
+  `make` if you don't have ninja installed).
 
 On the initial build, you can also speed things up with the environment
 variables `DEBUG` and `NO_CUDA`.
@@ -181,6 +272,8 @@ information for the code in `torch/csrc`. More information at:
 Python `setuptools` is pretty dumb, and always rebuilds every C file in a
 project.  If you install the ninja build system with `pip install ninja`,
 then PyTorch will use it to track dependencies correctly.
+If pytorch was already built, you will need to run `python setup.py clean` once
+after installing ninja for builds to succeed.
 
 #### Use CCache
 
@@ -247,9 +340,9 @@ than Linux, which are worth keeping in mind when fixing these problems.
 1. Symbols are NOT exported by default on Windows; instead, you have to explicitly
    mark a symbol as exported/imported in a header file with `__declspec(dllexport)` /
    `__declspec(dllimport)`.  We have codified this pattern into a set of macros
-   which follow the convention `*_API`, e.g., `AT_API` inside ATen. (Every separate
-   shared library needs a unique macro name, because symbol visibility is on a per
-   shared library basis.)
+   which follow the convention `*_API`, e.g., `CAFFE2_API` inside Caffe2 and ATen.
+   (Every separate shared library needs a unique macro name, because symbol visibility
+   is on a per shared library basis. See c10/macros/Macros.h for more details.)
 
    The upshot is if you see an "unresolved external" error in your Windows build, this
    is probably because you forgot to mark a function with `*_API`.  However, there is
@@ -266,8 +359,7 @@ than Linux, which are worth keeping in mind when fixing these problems.
 3. If you have a Windows box (we have a few on EC2 which you can request access to) and
    you want to run the build, the easiest way is to just run `.jenkins/pytorch/win-build.sh`.
    If you need to rebuild, run `REBUILD=1 .jenkins/pytorch/win-build.sh` (this will avoid
-   blowing away your Conda environment.)  I recommend opening `cmd.exe`, and then running
-   `bash` to work in a bash shell (which will make various Linux commands available.)
+   blowing away your Conda environment.)
 
 Even if you don't know anything about MSVC, you can use cmake to build simple programs on
 Windows; this can be helpful if you want to learn more about some peculiar linking behavior
@@ -293,6 +385,84 @@ cmake ..
 cmake --build .
 ```
 
+### Known MSVC (and MSVC with NVCC) bugs
+
+The PyTorch codebase sometimes likes to use exciting C++ features, and
+these exciting features lead to exciting bugs in Windows compilers.
+To add insult to injury, the error messages will often not tell you
+which line of code actually induced the erroring template instantiation.
+
+I've found the most effective way to debug these problems is to
+carefully read over diffs, keeping in mind known bugs in MSVC/NVCC.
+Here are a few well known pitfalls and workarounds:
+
+* This is not actually a bug per se, but in general, code generated by MSVC
+  is more sensitive to memory errors; you may have written some code
+  that does a use-after-free or stack overflows; on Linux the code
+  might work, but on Windows your program will crash.  ASAN may not
+  catch all of these problems: stay vigilant to the possibility that
+  your crash is due to a real memory problem.
+
+* (NVCC) `c10::optional` does not work when used from device code.  Don't use
+  it from kernels.  Upstream issue: https://github.com/akrzemi1/Optional/issues/58
+  and our local issue #10329.
+
+* `constexpr` generally works less well on MSVC.
+
+  * The idiom `static_assert(f() == f())` to test if `f` is constexpr
+    does not work; you'll get "error C2131: expression did not evaluate
+    to a constant".  Don't use these asserts on Windows.
+    (Example: `c10/util/intrusive_ptr.h`)
+
+* (NVCC) Code you access inside a `static_assert` will eagerly be
+  evaluated as if it were device code, and so you might get an error
+  that the code is "not accessible".
+
+```
+class A {
+  static A singleton_;
+  static constexpr inline A* singleton() {
+    return &singleton_;
+  }
+};
+static_assert(std::is_same(A*, decltype(A::singelton()))::value, "hmm");
+```
+
+* The compiler will run out of heap if you attempt to compile files that
+  are too large.  Splitting such files into separate files helps.
+  (Example: `THTensorMath`, `THTensorMoreMath`, `THTensorEvenMoreMath`.)
+
+### Running Clang-Tidy
+
+[Clang-Tidy](https://clang.llvm.org/extra/clang-tidy/index.html) is a C++
+linter and static analysis tool based on the clang compiler. We run clang-tidy
+in our CI to make sure that new C++ code is safe, sane and efficient. See our
+[.travis.yml](https://github.com/pytorch/pytorch/blob/master/.travis.yml) file
+for the simple commands we use for this.
+
+To run clang-tidy locally, follow these steps:
+
+1. Install clang-tidy. First, check if you already have clang-tidy by simply
+writing `clang-tidy` in your terminal. If you don't yet have clang-tidy, you
+should be able to install it easily with your package manager, e.g. by writing
+`apt-get install clang-tidy` on Ubuntu. See https://apt.llvm.org for details on
+how to install the latest version. Note that newer versions of clang-tidy will
+have more checks than older versions. In our CI, we run clang-tidy-6.0.
+
+2. Use our driver script to run clang-tidy over any changes relative to some
+   git revision (you may want to replace `HEAD~1` with `HEAD` to pick up
+   uncommitted changes). Changes are picked up based on a `git diff` with the
+   given revision:
+  ```sh
+  $ python tools/clang_tidy.py -d build -p torch/csrc --diff 'HEAD~1'
+  ```
+
+Above, it is assumed you are in the PyTorch root folder. `path/to/build` should
+be the path to where you built PyTorch from source, e.g. `build` in the PyTorch
+root folder if you used `setup.py build`. You can use `-c <clang-tidy-binary>`
+to change the clang-tidy this script uses. Make sure you have PyYaml installed,
+which is in PyTorch's `requirements.txt`.
+
 ## Caffe2 notes
 
 In 2018, we merged Caffe2 into the PyTorch source repository.  While the

README.md

@@ -15,17 +15,20 @@ We are in an early-release beta. Expect some adventures and rough edges.
   - [Binaries](#binaries)
   - [From Source](#from-source)
   - [Docker Image](#docker-image)
+  - [Building the Documentation](#building-the-documentation)
   - [Previous Versions](#previous-versions)
 - [Getting Started](#getting-started)
 - [Communication](#communication)
 - [Releases and Contributing](#releases-and-contributing)
 - [The Team](#the-team)
 
-| System | 2.7 | 3.5 |
-| --- | --- | --- |
-| Linux CPU | [![Build Status](https://ci.pytorch.org/jenkins/job/pytorch-master/badge/icon)](https://ci.pytorch.org/jenkins/job/pytorch-master/) | [![Build Status](https://ci.pytorch.org/jenkins/job/pytorch-master/badge/icon)](https://ci.pytorch.org/jenkins/job/pytorch-master/) |
-| Linux GPU | [![Build Status](https://ci.pytorch.org/jenkins/job/pytorch-master/badge/icon)](https://ci.pytorch.org/jenkins/job/pytorch-master/) | [![Build Status](https://ci.pytorch.org/jenkins/job/pytorch-master/badge/icon)](https://ci.pytorch.org/jenkins/job/pytorch-master/) |
-| Windows GPU | <center>—</center> | [![Build Status](https://ci.pytorch.org/jenkins/job/pytorch-builds/job/pytorch-win-ws2016-cuda9-cudnn7-py3-trigger/badge/icon)](https://ci.pytorch.org/jenkins/job/pytorch-builds/job/pytorch-win-ws2016-cuda9-cudnn7-py3-trigger/)
+| System | 2.7 | 3.5 | 3.6 |
+| :---: | :---: | :---: | :--: |
+| Linux CPU | [![Build Status](https://ci.pytorch.org/jenkins/job/pytorch-master/badge/icon)](https://ci.pytorch.org/jenkins/job/pytorch-master/) | [![Build Status](https://ci.pytorch.org/jenkins/job/pytorch-master/badge/icon)](https://ci.pytorch.org/jenkins/job/pytorch-master/) | <center>—</center> |
+| Linux GPU | [![Build Status](https://ci.pytorch.org/jenkins/job/pytorch-master/badge/icon)](https://ci.pytorch.org/jenkins/job/pytorch-master/) | [![Build Status](https://ci.pytorch.org/jenkins/job/pytorch-master/badge/icon)](https://ci.pytorch.org/jenkins/job/pytorch-master/) | <center>—</center> |
+| Windows GPU | <center>—</center> | [![Build Status](https://ci.pytorch.org/jenkins/job/pytorch-builds/job/pytorch-win-ws2016-cuda9-cudnn7-py3-trigger/badge/icon)](https://ci.pytorch.org/jenkins/job/pytorch-builds/job/pytorch-win-ws2016-cuda9-cudnn7-py3-trigger/) |  <center>—</center> |
+| Linux (ppc64le) CPU | [![Build Status](https://powerci.osuosl.org/job/pytorch-master-nightly-py2-linux-ppc64le/badge/icon)](https://powerci.osuosl.org/job/pytorch-master-nightly-py2-linux-ppc64le/) | — | [![Build Status](https://powerci.osuosl.org/job/pytorch-master-nightly-py3-linux-ppc64le/badge/icon)](https://powerci.osuosl.org/job/pytorch-master-nightly-py3-linux-ppc64le/) |
+| Linux (ppc64le) GPU | [![Build Status](https://powerci.osuosl.org/job/pytorch-linux-cuda9-cudnn7-py2-mpi-build-test-gpu/badge/icon)](https://powerci.osuosl.org/job/pytorch-linux-cuda9-cudnn7-py2-mpi-build-test-gpu/) | — | [![Build Status](https://powerci.osuosl.org/job/pytorch-linux-cuda92-cudnn7-py3-mpi-build-test-gpu/badge/icon)](https://powerci.osuosl.org/job/pytorch-linux-cuda92-cudnn7-py3-mpi-build-test-gpu/) |
 
 See also the [ci.pytorch.org HUD](https://ezyang.github.io/pytorch-ci-hud/build/pytorch-master).
 
@@ -56,8 +59,8 @@ If you use NumPy, then you have used Tensors (a.k.a ndarray).
 
 ![Tensor illustration](https://github.com/pytorch/pytorch/blob/master/docs/source/_static/img/tensor_illustration.png)
 
-PyTorch provides Tensors that can live either on the CPU or the GPU, and accelerate
-compute by a huge amount.
+PyTorch provides Tensors that can live either on the CPU or the GPU, and accelerates the
+computation by a huge amount.
 
 We provide a wide variety of tensor routines to accelerate and fit your scientific computation needs
 such as slicing, indexing, math operations, linear algebra, reductions.
@@ -76,7 +79,7 @@ change the way your network behaves arbitrarily with zero lag or overhead. Our i
 from several research papers on this topic, as well as current and past work such as
 [torch-autograd](https://github.com/twitter/torch-autograd),
 [autograd](https://github.com/HIPS/autograd),
-[Chainer](http://chainer.org), etc.
+[Chainer](https://chainer.org), etc.
 
 While this technique is not unique to PyTorch, it's one of the fastest implementations of it to date.
 You get the best of speed and flexibility for your crazy research.
@@ -87,7 +90,7 @@ You get the best of speed and flexibility for your crazy research.
 
 PyTorch is not a Python binding into a monolithic C++ framework.
 It is built to be deeply integrated into Python.
-You can use it naturally like you would use NumPy / SciPy / scikit-learn etc.
+You can use it naturally like you would use [NumPy](http://www.numpy.org/) / [SciPy](https://www.scipy.org/) / [scikit-learn](http://scikit-learn.org) etc.
 You can write your new neural network layers in Python itself, using your favorite libraries
 and use packages such as Cython and Numba.
 Our goal is to not reinvent the wheel where appropriate.
@@ -103,10 +106,9 @@ We hope you never spend hours debugging your code because of bad stack traces or
 ### Fast and Lean
 
 PyTorch has minimal framework overhead. We integrate acceleration libraries
-such as Intel MKL and NVIDIA (cuDNN, NCCL) to maximize speed.
+such as [Intel MKL](https://software.intel.com/mkl) and NVIDIA (cuDNN, NCCL) to maximize speed.
 At the core, its CPU and GPU Tensor and neural network backends
-(TH, THC, THNN, THCUNN) are written as independent libraries with a C99 API.
-They are mature and have been tested for years.
+(TH, THC, THNN, THCUNN) are mature and have been tested for years.
 
 Hence, PyTorch is quite fast – whether you run small or large neural networks.
 
@@ -121,10 +123,10 @@ Writing new neural network modules, or interfacing with PyTorch's Tensor API was
 and with minimal abstractions.
 
 You can write new neural network layers in Python using the torch API
-[or your favorite NumPy-based libraries such as SciPy](http://pytorch.org/tutorials/advanced/numpy_extensions_tutorial.html).
+[or your favorite NumPy-based libraries such as SciPy](https://pytorch.org/tutorials/advanced/numpy_extensions_tutorial.html).
 
 If you want to write your layers in C/C++, we provide a convenient extension API that is efficient and with minimal boilerplate.
-There is no wrapper code that needs to be written. You can see [a tutorial here](http://pytorch.org/tutorials/advanced/cpp_extension.html) and [an example here](https://github.com/pytorch/extension-cpp).
+There is no wrapper code that needs to be written. You can see [a tutorial here](https://pytorch.org/tutorials/advanced/cpp_extension.html) and [an example here](https://github.com/pytorch/extension-cpp).
 
 
 ## Installation
@@ -132,7 +134,7 @@ There is no wrapper code that needs to be written. You can see [a tutorial here]
 ### Binaries
 Commands to install from binaries via Conda or pip wheels are on our website:
 
-[http://pytorch.org](http://pytorch.org)
+[https://pytorch.org](https://pytorch.org)
 
 ### From Source
 
@@ -148,7 +150,9 @@ If you want to compile with CUDA support, install
 If you want to disable CUDA support, export environment variable `NO_CUDA=1`.
 Other potentially useful environment variables may be found in `setup.py`.
 
-If you want to build on Windows, Visual Studio 2017 and NVTX are also needed.
+If you want to build on Windows, Visual Studio 2017 14.11 toolset and NVTX are also needed.
+Especially, for CUDA 8 build on Windows, there will be an additional requirement for VS 2015 Update 3 and a patch for it.
+The details of the patch can be found out [here](https://support.microsoft.com/en-gb/help/4020481/fix-link-exe-crashes-with-a-fatal-lnk1000-error-when-you-use-wholearch).
 
 #### Install optional dependencies
 
@@ -161,7 +165,7 @@ conda install numpy pyyaml mkl mkl-include setuptools cmake cffi typing
 conda install -c mingfeima mkldnn
 
 # Add LAPACK support for the GPU
-conda install -c pytorch magma-cuda80 # or magma-cuda90 if CUDA 9
+conda install -c pytorch magma-cuda92 # or [magma-cuda80 | magma-cuda91] depending on your cuda version
 ```
 
 On macOS
@@ -196,11 +200,11 @@ On Windows
 set "VS150COMNTOOLS=C:\Program Files (x86)\Microsoft Visual Studio\2017\Enterprise\VC\Auxiliary\Build"
 set CMAKE_GENERATOR=Visual Studio 15 2017 Win64
 set DISTUTILS_USE_SDK=1
-REM The following line is needed for Python 2.7, but the support for it is very experimental.
+REM The following two lines are needed for Python 2.7, but the support for it is very experimental.
 set MSSdk=1
-REM As for CUDA 8, VS2015 Update 2 or up is required to build PyTorch. Use the following two lines.
-set "PREBUILD_COMMAND=%VS140COMNTOOLS%\..\..\VC\vcvarsall.bat"
-set PREBUILD_COMMAND_ARGS=x64
+set FORCE_PY27_BUILD=1
+REM As for CUDA 8, VS2015 Update 3 is also required to build PyTorch. Use the following line.
+set "CUDAHOSTCXX=%VS140COMNTOOLS%\..\..\VC\bin\amd64\cl.exe"
 
 call "%VS150COMNTOOLS%\vcvarsall.bat" x64 -vcvars_ver=14.11
 python setup.py install
@@ -208,7 +212,7 @@ python setup.py install
 
 ### Docker image
 
-Dockerfile is supplied to build images with cuda support and cudnn v7. Build as usual
+Dockerfile is supplied to build images with cuda support and cudnn v7. You can pass `-e PYTHON_VERSION=x.y` flag to specify which python version is to be used by Miniconda, or leave it unset to use the default. Build as usual
 ```
 docker build -t pytorch -f docker/pytorch/Dockerfile .
 ```
@@ -222,24 +226,36 @@ Please note that PyTorch uses shared memory to share data between processes, so
 for multithreaded data loaders) the default shared memory segment size that container runs with is not enough, and you
 should increase shared memory size either with `--ipc=host` or `--shm-size` command line options to `nvidia-docker run`.
 
+### Building the Documentation
+
+To build documentation in various formats, you will need [Sphinx](http://www.sphinx-doc.org) and the
+readthedocs theme.
+
+```
+cd docs/
+pip install -r requirements.txt
+```
+You can then build the documentation by running ``make <format>`` from the
+``docs/`` folder. Run ``make`` to get a list of all available output formats.
+
 ### Previous Versions
 
 Installation instructions and binaries for previous PyTorch versions may be found
-on [our website](http://pytorch.org/previous-versions/).
+on [our website](https://pytorch.org/previous-versions).
 
 
 ## Getting Started
 
 Three pointers to get you started:
-- [Tutorials: get you started with understanding and using PyTorch](http://pytorch.org/tutorials/)
+- [Tutorials: get you started with understanding and using PyTorch](https://pytorch.org/tutorials/)
 - [Examples: easy to understand pytorch code across all domains](https://github.com/pytorch/examples)
-- [The API Reference](http://pytorch.org/docs/)
+- [The API Reference](https://pytorch.org/docs/)
 
 ## Communication
-* forums: discuss implementations, research, etc. http://discuss.pytorch.org
+* forums: discuss implementations, research, etc. https://discuss.pytorch.org
 * GitHub issues: bug reports, feature requests, install issues, RFCs, thoughts, etc.
 * Slack: general chat, online discussions, collaboration etc. https://pytorch.slack.com/ . Our slack channel is invite-only to promote a healthy balance between power-users and beginners. If you need a slack invite, ping us at slack@pytorch.org
-* newsletter: no-noise, one-way email newsletter with important announcements about pytorch. You can sign-up here: http://eepurl.com/cbG0rv
+* newsletter: no-noise, one-way email newsletter with important announcements about pytorch. You can sign-up here: https://eepurl.com/cbG0rv
 
 ## Releases and Contributing
 
@@ -259,3 +275,7 @@ PyTorch is currently maintained by [Adam Paszke](https://apaszke.github.io/), [S
 A non-exhaustive but growing list needs to mention: Trevor Killeen, Sasank Chilamkurthy, Sergey Zagoruyko, Adam Lerer, Francisco Massa, Alykhan Tejani, Luca Antiga, Alban Desmaison, Andreas Kopf, James Bradbury, Zeming Lin, Yuandong Tian, Guillaume Lample, Marat Dukhan, Natalia Gimelshein, Christian Sarofeen, Martin Raison, Edward Yang, Zachary Devito.
 
 Note: this project is unrelated to [hughperkins/pytorch](https://github.com/hughperkins/pytorch) with the same name. Hugh is a valuable contributor in the Torch community and has helped with many things Torch and PyTorch.
+
+## License
+
+PyTorch is BSD-style licensed, as found in the LICENSE file.

aten/.flake8

@@ -1,3 +0,0 @@
-[flake8]
-max-line-length = 120
-

aten/.gitignore

@@ -1,3 +0,0 @@
-__pycache__/
-build/
-*.pyc

aten/CMakeLists.txt

@@ -1,22 +1,5 @@
-if (CAFFE2_CMAKE_BUILDING_WITH_MAIN_REPO)
-  if (NOT BUILD_ATEN)
-    return()
-  endif()
-else()
-  cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
-  project(ATen CXX C)
-  include(CMakeDependentOption)
-  option(USE_CUDA "Use CUDA" ON)
-  option(USE_ROCM "Use ROCm" OFF)
-  option(USE_CUDNN "Use cuDNN" ON)
-  option(USE_MKLDNN "Use MKLDNN" ON)
-  cmake_dependent_option(
-      USE_CUDNN "Use cuDNN" ON
-      "USE_CUDA" OFF)
-  option(ATEN_NO_TEST "Do not build ATen test binaries" OFF)
-
-  # Flag for shared dependencies
-  set(BUILD_ATEN ON)
+if (BUILD_ATEN_MOBILE)
+  return()
 endif()
 
 # Find modules
@@ -45,32 +28,6 @@ SET(ATEN_INSTALL_BIN_SUBDIR "bin" CACHE PATH "ATen install binary subdirectory")
 SET(ATEN_INSTALL_LIB_SUBDIR "lib" CACHE PATH "ATen install library subdirectory")
 SET(ATEN_INSTALL_INCLUDE_SUBDIR "include" CACHE PATH "ATen install include subdirectory")
 
-if (NOT CAFFE2_CMAKE_BUILDING_WITH_MAIN_REPO)
-  # ---[ Build variables set within the cmake tree
-  include(../cmake/BuildVariables.cmake)
-  set(CAFFE2_WHITELIST "" CACHE STRING "A whitelist file of files that one should build.")
-
-  # ---[ Misc checks to cope with various compiler modes
-  include(../cmake/MiscCheck.cmake)
-
-  # External projects
-  include(ExternalProject)
-
-  # ---[ Utils
-  # TODO: merge the following 3 files into cmake/public/utils.cmake.
-  include(../cmake/Utils.cmake)
-  include(../cmake/public/utils.cmake)
-
-  # ---[ Dependencies
-  include(../cmake/Dependencies.cmake)
-  list(APPEND ATen_CPU_INCLUDE ${Caffe2_CPU_INCLUDE})
-  list(APPEND ATen_CUDA_INCLUDE ${Caffe2_GPU_INCLUDE})
-  list(APPEND ATen_CPU_DEPENDENCY_LIBS ${Caffe2_DEPENDENCY_LIBS})
-  list(APPEND ATen_CUDA_DEPENDENCY_LIBS ${Caffe2_CUDA_DEPENDENCY_LIBS})
-  list(APPEND ATen_PUBLIC_CUDA_DEPENDENCY_LIBS
-    ${Caffe2_PUBLIC_CUDA_DEPENDENCY_LIBS})
-endif()
-
 if(USE_CUDA)
   list(APPEND ATen_CUDA_INCLUDE ${CUDA_INCLUDE_DIRS})
 endif()
@@ -80,14 +37,20 @@ add_subdirectory(src/TH)
 set(TH_CPU_INCLUDE
   # dense
   ${CMAKE_CURRENT_SOURCE_DIR}/src/TH
-  ${CMAKE_CURRENT_SOURCE_DIR}/src/THC
   ${CMAKE_CURRENT_BINARY_DIR}/src/TH
-  ${CMAKE_CURRENT_BINARY_DIR}/src/THC
-
   ${CMAKE_CURRENT_SOURCE_DIR}/src
   ${CMAKE_CURRENT_BINARY_DIR}/src
   ${CMAKE_BINARY_DIR}/aten/src)
 list(APPEND ATen_CPU_INCLUDE ${TH_CPU_INCLUDE})
+
+if(USE_CUDA OR USE_ROCM)
+  set(TH_CUDA_INCLUDE
+    # dense
+    ${CMAKE_CURRENT_SOURCE_DIR}/src/THC
+    ${CMAKE_CURRENT_BINARY_DIR}/src/THC)
+  list(APPEND ATen_CUDA_INCLUDE ${TH_CUDA_INCLUDE})
+endif()
+
 add_subdirectory(src/THNN)
 
 # Find the HIP package, set the HIP paths, load the HIP CMake.
@@ -129,15 +92,14 @@ list(APPEND ATen_CPU_INCLUDE
   ${CMAKE_CURRENT_BINARY_DIR}/src/ATen)
 add_subdirectory(src/ATen)
 
-if (CAFFE2_CMAKE_BUILDING_WITH_MAIN_REPO)
-  # Pass source, includes, and libs to parent
-  set(ATen_CPU_SRCS ${ATen_CPU_SRCS} PARENT_SCOPE)
-  set(ATen_CUDA_SRCS ${ATen_CUDA_SRCS} PARENT_SCOPE)
-  set(ATen_CPU_TEST_SRCS ${ATen_CPU_TEST_SRCS} PARENT_SCOPE)
-  set(ATen_CUDA_TEST_SRCS ${ATen_CUDA_TEST_SRCS} PARENT_SCOPE)
-  set(ATen_CPU_INCLUDE ${ATen_CPU_INCLUDE} PARENT_SCOPE)
-  set(ATen_CUDA_INCLUDE ${ATen_CUDA_INCLUDE} PARENT_SCOPE)
-  set(ATen_THIRD_PARTY_INCLUDE ${ATen_THIRD_PARTY_INCLUDE} PARENT_SCOPE)
-  set(ATen_CPU_DEPENDENCY_LIBS ${ATen_CPU_DEPENDENCY_LIBS} PARENT_SCOPE)
-  set(ATen_CUDA_DEPENDENCY_LIBS ${ATen_CUDA_DEPENDENCY_LIBS} PARENT_SCOPE)
-endif()
+# Pass source, includes, and libs to parent
+set(ATen_CPU_SRCS ${ATen_CPU_SRCS} PARENT_SCOPE)
+set(ATen_CUDA_SRCS ${ATen_CUDA_SRCS} PARENT_SCOPE)
+set(ATen_CPU_TEST_SRCS ${ATen_CPU_TEST_SRCS} PARENT_SCOPE)
+set(ATen_CUDA_TEST_SRCS ${ATen_CUDA_TEST_SRCS} PARENT_SCOPE)
+set(ATen_CPU_INCLUDE ${ATen_CPU_INCLUDE} PARENT_SCOPE)
+set(ATen_CUDA_INCLUDE ${ATen_CUDA_INCLUDE} PARENT_SCOPE)
+set(ATen_THIRD_PARTY_INCLUDE ${ATen_THIRD_PARTY_INCLUDE} PARENT_SCOPE)
+set(ATen_CPU_DEPENDENCY_LIBS ${ATen_CPU_DEPENDENCY_LIBS} PARENT_SCOPE)
+set(ATen_CUDA_DEPENDENCY_LIBS ${ATen_CUDA_DEPENDENCY_LIBS} PARENT_SCOPE)
+set(ATen_CORE_TEST_SRCS ${ATen_CORE_TEST_SRCS} PARENT_SCOPE)

aten/README.md

@@ -1,258 +0,0 @@
-# ATen: A TENsor library
-
-ATen is a simple tensor library thats exposes the Tensor operations in Torch
-and PyTorch directly in C++11. The wrapper respects the semantics of operators
-in PyTorch, except minor details due to differences between C++ and Python in
-the way default arguments are handled. See the [documentation for tensors](http://pytorch.org/docs/tensors.html) in PyTorch for what these operations do.
-ATen's API is auto-generated from the same declarations PyTorch uses so the
-two APIs will track each other over time.
-
-Tensor types are resolved dynamically, such that the API is generic and
-does not include templates. That is, there is one `Tensor` type. It can hold a
-CPU or CUDA Tensor, and the tensor may have Doubles, Float, Ints, etc. This design
-makes it easy to write generic code without templating everything.
-
-See the _generated_ [`Tensor.h` file](doc/Tensor.h) and [`Functions.h` file](doc/Functions.h) for the provided API. Excerpt:
-```c++
-Tensor atan2(const Tensor & other) const;
-Tensor & atan2_(const Tensor & other);
-Tensor pow(Scalar exponent) const;
-Tensor pow(const Tensor & exponent) const;
-Tensor & pow_(Scalar exponent);
-Tensor & pow_(const Tensor & exponent);
-Tensor lerp(const Tensor & end, Scalar weight) const;
-Tensor & lerp_(const Tensor & end, Scalar weight);
-Tensor histc() const;
-Tensor histc(int64_t bins) const;
-Tensor histc(int64_t bins, Scalar min) const;
-Tensor histc(int64_t bins, Scalar min, Scalar max) const;
-```
-
-Inplace operations are also provided, and always suffixed by `_` to indicate they will modify the Tensor.
-
-### Installation
-
-TH/THC/THNN/THCUNN are provided (as git subtrees), so the repo is standalone. You will need a C++11 compiler, cmake, and the pyyaml python package.
-```
-
-# Install pyyaml used by python code generation to read API declarations
-
-# macOS: if you don't have pip
-sudo easy_install pip
-# Ubuntu: if you don't have pip
-apt-get -y install python-pip
-
-# if you don't have pyyaml
-sudo pip install pyyaml
-
-mkdir build
-cd build
-cmake .. -DCMAKE_INSTALL_PREFIX=/where/you/want # specify your dest directory
-# cmake .. -DUSE_NVRTC=ON -DUSE_TENSORRT=OFF -DCMAKE_INSTALL_PREFIX=../install -DCAFFE2_CMAKE_BUILDING_WITH_MAIN_REPO=OFF -DUSE_CUDA=ON # for CUDA
-# cmake .. -DUSE_CUDA=OFF  # for CPU only machines
-make install
-```
-
-### Example usage
-
-Here is a simple example; again, the syntax follows Torch semantics.
-
-```c++
-using namespace at; // assumed in the following
-
-Tensor d = CPU(kFloat).ones({3, 4});
-Tensor r = CPU(kFloat).zeros({3,4});
-for(auto i = 0; i < 100000; i++) {
-  r = r.add(d);
-  // equivalently
-  r = r + d;
-  // or
-  r += d;
-}
-```
-
-Want this running on the GPU?
-```c++
-using namespace at; // assumed in the following
-
-Tensor d = CUDA(kFloat).ones({3, 4});
-Tensor r = CUDA(kFloat).zeros({3,4});
-for(auto i = 0; i < 100000; i++) {
-  r = r.add(d);
-  // equivalently
-  r = r + d;
-  // or
-  r += d;
-}
-```
-
-Expressions like `CUDA(kFloat)` are first-class `at::Type` objects that represent
-the type of a Tensor and are used to create Tensors when their type cannot be
-inferred. See the _generated_ [Type header](doc/Type.h) for its API.
-
-See more in [sample files](src/ATen/test).
-
-### Creating your kernel
-
-It is easy to create new kernels, thanks to the `dispatch<>()` templated function. Example:
-```c++
-
-// a simple sum kernel (for CPU only)
-template<typename T>
-struct sum_op {
-  // dispatch handles variable arguments for you
-  Tensor CPU(const Type & t, Tensor & x_)
-  {
-    Tensor x = x_.contiguous();
-    auto x_p = x.data<T>();
-    int64_t size = x.numel();
-    T sum = 0;
-    for(int64_t i = 0; i < size; i++) {
-      sum += x_p[i];
-    }
-    return sum;
-  };
-  Tensor CUDA(Tensor& x) {
-    throw std::invalid_argument("device not supported");
-  };
-};
-
-Tensor a = CPU(kFloat).rand({3, 7});
-std::cout << a << std::endl;
-std::cout << dispatch<sum_op>(a.type(),a) << " == " << a.sum() << std::endl;
-```
-
-### Efficient access to tensor elements
-
-When using Tensor-wide operations, the relative cost of dynamic dispatch is very small.
-However, there are cases, especially in your own kernels, where efficient element-wise access is needed,
-and the cost of dynamic dispatch inside the element-wise loop is very high.
-ATen provides _accessors_ that are created with a single dynamic check that a Tensor is the type and number of
-dimensions. Accessors then expose an API for accessing the Tensor elements efficiently:
-
-```c++
-
-Tensor foo = CPU(kFloat).rand({12,12});
-
-// assert foo is 2-dimensional and holds floats.
-auto foo_a = foo.accessor<float,2>();
-float trace = 0;
-
-for(int i = 0; i < foo_a.size(0); i++) {
-  // use the accessor foo_a to get tensor data.
-  trace += foo_a[i][i];
-}
-```
-
-Accessors are temporary views of a Tensor. They are only valid for the lifetime of the tensor that they
-view and hence should only be used locally in a function, like iterators.
-
-### Using externally created data
-
-If you already have your tensor data allocated in memory (CPU or CUDA),
-you can view that memory as a Tensor in ATen:
-
-```c++
-float data[] = { 1, 2, 3,
-                 4, 5, 6};
-auto f = CPU(kFloat).tensorFromBlob(data, {2,3});
-cout << f << endl;
-```
-
-These tensors cannot be resized because ATen does not own the memory, but otherwise
-behave as normal tensors.
-
-### Scalars and zero-dimensional tensors
-
-In addition to the `Tensor` objects, ATen also includes `Scalar`s that represent a single number.
-Like a Tensor, Scalars are dynamically typed and can hold any one of ATen's [number types](doc/Type.h).
-Scalars can be implicitly constructed from C++ number types. Scalars are needed because some functions like `addmm` take numbers along with Tensors and expect these
-numbers to be the same dynamic type as the tensor. They are also used in the API to indicate places where
-a function will _always_ return a Scalar value, like `sum`.
-
-```c++
-Tensor addmm(Scalar beta, const Tensor & self,
-             Scalar alpha, const Tensor & mat1,
-             const Tensor & mat2);
-Scalar sum(const Tensor & self);
-
-//usage
-Tensor a = ...
-Tensor b = ...
-Tensor c = ...
-Tensor r = addmm(1.0, a, .5, b, c);
-```
-
-In addition to Scalars, ATen also allows Tensor objects to be zero-dimensional. These Tensors hold
-a single value and they can be references to a single element in a larger Tensor. They can be used anywhere a Tensor is expected. They are normally created by operators like `select` which reduce the dimensions of
-a Tensor.
-
-```c++
-Tensor two = CPU(kFloat).rand({10,20});
-two[1][2] = 4;
-//~~~~~~~  zero-dimensional Tensor
-```
-
-It is possible to convert between Scalar and zero-dim Tensors:
-
-```c++
-Tensor zero_dim = CPU(kFloat).scalarTensor(4);
-Scalar from_tensor = Scalar(zero_dim); //only valid when zero_dim.dim() == 0;
-```
-
-### Avoiding unnecessary CUDA synchronization in your kernels when using Scalars
-
-Moving a single number from the GPU to the CPU introduces a synchronization point
-that can add latency to your program. In certain cases the result of a GPU operator like `sum` which
-returns a Scalar may be plugged into another GPU operator as an argument. If Scalars were always copied
-to the CPU, this would result in 2 copies. To avoid these synchronizations, Scalar objects can be
-optionally backed by a zero-dim Tensor, and are only copied to the CPU when requested.
-
-```c++
-auto a = CUDA(kFloat).rand({3,4});
-Scalar on_gpu = Scalar(a[1][1]); //backed by zero-dim Tensor
-assert(on_gpu.isBackedByTensor());
-
-double value = on_gpu.toDouble(); // copied to CPU, if it was backed by GPU Tensor.
-Scalar svalue = on_gpu.local(); // force the Scalar to become local to CPU.
-
-// get the scalar as a zero-dim tensor. If it was already backed
-// by a zero-dim Tensor then this op has no synchronization.
-// if the Scalar was local on CPU, it performs the copy
-Tensor same_tensor = CUDA(kFloat).scalarTensor(on_gpu);
-```
-
-Operators aware of the location of Scalars can arrange to do the minimal number of copies required.
-
-### Developer notes
-
-ATen relies heavily on code generation to automatically generate headers
-and implementations for all of the tensor methods it supports.  The main
-entry point for the script which does all this work is
-[`src/ATen/gen.py`](src/ATen/gen.py), which ingests
-[`src/ATen/Declarations.cwrap`](src/ATen/Declarations.cwrap),
-[`src/ATen/nn.yaml`](src/ATen/nn.yaml),
-[`src/ATen/native/native_functions.yaml`](src/ATen/native/native_functions.yaml) and the THNN/THCUNN headers and
-produces all of the headers and wrapping code necessary to generate
-the ATen interface.
-
-If you need to understand how ATen understands a declaration after all
-of this processing occurs, it's helpful to look at the generated file
-`Declarations.yaml` (NB: not cwrap) which contains information for all
-ATen methods in a uniform manner.  This file is utilized by PyTorch
-which further extends the ATen interface with support for automatic
-differentation.
-
-#### Note [ATen preprocessor philosophy]
-
-ATen is designed to be simple to use, and one of the things this implies is
-that it should not be necessary to use preprocessor macros when using ATen;
-we would rather provide all symbols, even for functionality that is not
-available on the system ATen is running on.
-
-This means that internally inside ATen, whereas other libraries might
-simply omit source files for, e.g., CuDNN, when CuDNN libraries are not
-installed, ATen will always build these source files, compiling stub
-functions for anything that is not available.  ATen never uses
-`AT_ENABLED_CUDA()` in header files, and all types in ATen's public API
-are always available no matter your build configuration.

aten/src/ATen/ATen.h

@@ -1,23 +1,24 @@
 #pragma once
 
-#include "ATen/ATenGeneral.h"
-#include "ATen/CPUGeneral.h"
 #include "ATen/Allocator.h"
-#include "ATen/Scalar.h"
-#include "ATen/Type.h"
-#include "ATen/Generator.h"
+#include "ATen/CPUGeneral.h"
 #include "ATen/Context.h"
-#include "ATen/Storage.h"
-#include "ATen/Tensor.h"
 #include "ATen/Device.h"
-#include "ATen/TensorGeometry.h"
-#include "ATen/Functions.h"
-#include "ATen/Formatting.h"
-#include "ATen/TensorOperators.h"
-#include "ATen/TensorMethods.h"
-#include "ATen/Dispatch.h"
-#include "ATen/DimVector.h"
 #include "ATen/DeviceGuard.h"
-#include "ATen/TensorOptions.h"
-#include "ATen/Layout.h"
-#include "ATen/OptionsGuard.h"
+#include "ATen/DimVector.h"
+#include "ATen/Dispatch.h"
+#include "ATen/Formatting.h"
+#include "ATen/Functions.h"
+#include "ATen/ScalarOps.h"
+#include "ATen/Tensor.h"
+#include "ATen/TensorGeometry.h"
+#include "ATen/TensorOperators.h"
+#include "ATen/Type.h"
+#include "ATen/core/ATenGeneral.h"
+#include "ATen/core/Generator.h"
+#include <c10/core/Layout.h>
+#include "ATen/core/Scalar.h"
+#include <c10/core/Storage.h>
+#include "ATen/core/TensorMethods.h"
+#include "ATen/core/TensorOptions.h"
+#include <c10/util/Exception.h>

aten/src/ATen/ATenGeneral.h

@@ -1,11 +0,0 @@
-#pragma once
-
-#ifdef _WIN32
-# if defined(ATen_cpu_EXPORTS) || defined(caffe2_EXPORTS)
-#  define AT_API __declspec(dllexport)
-# else
-#  define AT_API __declspec(dllimport)
-# endif
-#else
-# define AT_API
-#endif

aten/src/ATen/AccumulateType.h

@@ -1,14 +1,17 @@
 #pragma once
 #include "ATen/Config.h"
-#include "ATen/Half.h"
+#include "ATen/core/Half.h"
 
 // Defines the accumulation type for a scalar type.
 // Example:
 //   using accscalar_t = acc_type<scalar_t, true>;
 
-#ifdef __CUDACC__
+#if defined(__CUDACC__)
 #include <cuda.h>
 #include <cuda_fp16.h>
+#elif defined(__HIPCC__)
+#include <hip/hip_runtime.h>
+#include <hip/hip_fp16.h>
 #endif
 
 namespace at {
@@ -16,7 +19,7 @@ namespace at {
 template <typename T, bool is_cuda>
 struct AccumulateType { };
 
-#ifdef __CUDACC__
+#if defined(__CUDACC__) || defined(__HIPCC__)
 template <> struct AccumulateType<half, true> { using type = float; };
 #endif
 template <> struct AccumulateType<Half, true> { using type = float; };

aten/src/ATen/Allocator.cpp

@@ -1,14 +0,0 @@
-#include <ATen/Allocator.h>
-
-namespace at {
-
-static void deleteInefficientStdFunctionContext(void* ptr) {
-  delete static_cast<InefficientStdFunctionContext*>(ptr);
-}
-
-at::DataPtr
-InefficientStdFunctionContext::makeDataPtr(void* ptr, const std::function<void(void*)>& deleter, Device device) {
-  return {ptr, new InefficientStdFunctionContext({ptr, deleter}), &deleteInefficientStdFunctionContext, device};
-}
-
-} // namespace at

aten/src/ATen/Allocator.h

@@ -1,101 +1,2 @@
 #pragma once
-
-#include <memory>
-#include <stddef.h>
-
-#include <ATen/Error.h>
-#include <ATen/Retainable.h>
-#include <ATen/Device.h>
-#include <ATen/detail/UniqueVoidPtr.h>
-
-namespace at {
-
-// A DataPtr is a unique pointer (with an attached deleter and some
-// context for the deleter) to some memory, which also records what
-// device is for its data.
-//
-// nullptr DataPtrs can still have a nontrivial device; this allows
-// us to treat zero-size allocations uniformly with non-zero allocations.
-//
-class DataPtr {
-private:
-  detail::UniqueVoidPtr ptr_;
-  Device device_;
-public:
-  // Choice of CPU here is arbitrary; if there's an "undefined" device
-  // we could use that too
-  DataPtr() : ptr_(), device_(kCPU) {}
-  DataPtr(void* data, Device device)
-    : ptr_(data), device_(device) {}
-  DataPtr(void* data, void* ctx, DeleterFnPtr ctx_deleter, Device device)
-    : ptr_(data, ctx, ctx_deleter), device_(device) {}
-  void* operator->() const { return ptr_.get(); }
-  void* get() const { return ptr_.get(); }
-  void* get_context() const { return ptr_.get_context(); }
-  void* release_context() { return ptr_.release_context(); }
-  operator bool() const { return static_cast<bool>(ptr_); }
-  template <typename T>
-  T* cast_context(DeleterFnPtr expected_deleter) const {
-    return ptr_.cast_context<T>(expected_deleter);
-  }
-  Device device() const { return device_; }
-};
-
-// NB: Device is NOT tested for here; a CUDA nullptr is as much a nullptr as a
-// CPU nullptr
-
-inline bool operator==(const at::DataPtr& dp, std::nullptr_t) noexcept { return !dp; }
-inline bool operator==(std::nullptr_t, const at::DataPtr& dp) noexcept { return !dp; }
-inline bool operator!=(const at::DataPtr& dp, std::nullptr_t) noexcept { return dp; }
-inline bool operator!=(std::nullptr_t, const at::DataPtr& dp) noexcept { return dp; }
-
-// Note [raw_allocate/raw_deallocate and Thrust]
-// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-// Thrust's support for custom allocators requires us to write something
-// like this:
-//
-//  class ThrustAllocator {
-//    char* allocate(size_t);
-//    void deallocate(char*, size_t);
-//  };
-//
-// This is not good for our unique_ptr based allocator interface, as
-// there is no way to get to the context when we free.
-//
-// However, in some cases the context is exactly the same as
-// the data pointer.  In this case, we can support the "raw"
-// allocate and deallocate interface.  This is what
-// raw_deleter signifies.  By default, it returns a nullptr, which means that
-// the raw interface is not implemented.  Be sure to implement it whenever
-// possible, or the raw interface will incorrectly reported as unsupported,
-// when it is actually possible.
-
-struct Allocator {
-  virtual ~Allocator() {}
-  virtual at::DataPtr allocate(size_t n) const = 0;
-
-  // If this returns a non nullptr, it means that allocate()
-  // is guaranteed to return a unique_ptr with this deleter attached;
-  // it means the rawAllocate and rawDeallocate APIs are safe to use.
-  // This function MUST always return the same BoundDeleter.
-  virtual DeleterFnPtr raw_deleter() const { return nullptr; }
-  void* raw_allocate(size_t n) {
-    auto dptr = allocate(n);
-    AT_ASSERT(dptr.get() == dptr.get_context());
-    return dptr.release_context();
-  }
-  void raw_deallocate(void* ptr) {
-    auto d = raw_deleter();
-    AT_ASSERT(d);
-    d(ptr);
-  }
-};
-
-struct AT_API InefficientStdFunctionContext {
-  std::unique_ptr<void, std::function<void(void*)>> ptr_;
-  InefficientStdFunctionContext(std::unique_ptr<void, std::function<void(void*)>>&& ptr)
-    : ptr_(std::move(ptr)) {}
-  static at::DataPtr makeDataPtr(void* ptr, const std::function<void(void*)>& deleter, Device device);
-};
-
-}  // namespace at
+#include <c10/core/Allocator.h>

aten/src/ATen/ArrayRef.h

@@ -1,192 +1,2 @@
-//===--- ArrayRef.h - Array Reference Wrapper -------------------*- C++ -*-===//
-//
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
-//
-//===----------------------------------------------------------------------===//
-
-// ATen: modified from llvm::ArrayRef.
-// removed llvm-specific functionality
-// removed some implicit const -> non-const conversions that rely on
-// complicated std::enable_if meta-programming
-// removed a bunch of slice variants for simplicity...
-
 #pragma once
-
-#include <ATen/Error.h>
-#include <ATen/SmallVector.h>
-
-#include <array>
-#include <iterator>
-#include <vector>
-
-namespace at {
-  /// ArrayRef - Represent a constant reference to an array (0 or more elements
-  /// consecutively in memory), i.e. a start pointer and a length.  It allows
-  /// various APIs to take consecutive elements easily and conveniently.
-  ///
-  /// This class does not own the underlying data, it is expected to be used in
-  /// situations where the data resides in some other buffer, whose lifetime
-  /// extends past that of the ArrayRef. For this reason, it is not in general
-  /// safe to store an ArrayRef.
-  ///
-  /// This is intended to be trivially copyable, so it should be passed by
-  /// value.
-  template<typename T>
-  class ArrayRef {
-  public:
-    typedef const T *iterator;
-    typedef const T *const_iterator;
-    typedef size_t size_type;
-
-    typedef std::reverse_iterator<iterator> reverse_iterator;
-
-  private:
-    /// The start of the array, in an external buffer.
-    const T *Data;
-
-    /// The number of elements.
-    size_type Length;
-
-  public:
-    /// @name Constructors
-    /// @{
-
-    /// Construct an empty ArrayRef.
-    /*implicit*/ ArrayRef() : Data(nullptr), Length(0) {}
-
-    /// Construct an ArrayRef from a single element.
-    /*implicit*/ ArrayRef(const T &OneElt)
-      : Data(&OneElt), Length(1) {}
-
-    /// Construct an ArrayRef from a pointer and length.
-    /*implicit*/ ArrayRef(const T *data, size_t length)
-      : Data(data), Length(length) {}
-
-    /// Construct an ArrayRef from a range.
-    ArrayRef(const T *begin, const T *end)
-      : Data(begin), Length(end - begin) {}
-
-    /// Construct an ArrayRef from a SmallVector. This is templated in order to
-    /// avoid instantiating SmallVectorTemplateCommon<T> whenever we
-    /// copy-construct an ArrayRef.
-    template<typename U>
-    /*implicit*/ ArrayRef(const SmallVectorTemplateCommon<T, U> &Vec)
-      : Data(Vec.data()), Length(Vec.size()) {
-    }
-
-    /// Construct an ArrayRef from a std::vector.
-    template<typename A>
-    /*implicit*/ ArrayRef(const std::vector<T, A> &Vec)
-      : Data(Vec.data()), Length(Vec.size()) {}
-
-    /// Construct an ArrayRef from a std::array
-    template <size_t N>
-    /*implicit*/ constexpr ArrayRef(const std::array<T, N> &Arr)
-        : Data(Arr.data()), Length(N) {}
-
-    /// Construct an ArrayRef from a C array.
-    template <size_t N>
-    /*implicit*/ constexpr ArrayRef(const T (&Arr)[N]) : Data(Arr), Length(N) {}
-
-    /// Construct an ArrayRef from a std::initializer_list.
-    /*implicit*/ ArrayRef(const std::initializer_list<T> &Vec)
-    : Data(Vec.begin() == Vec.end() ? (T*)nullptr : Vec.begin()),
-      Length(Vec.size()) {}
-
-    /// @}
-    /// @name Simple Operations
-    /// @{
-
-    const_iterator begin() const { return Data; }
-    const_iterator end() const { return Data + Length; }
-
-    reverse_iterator rbegin() const { return reverse_iterator(end()); }
-    reverse_iterator rend() const { return reverse_iterator(begin()); }
-
-    /// empty - Check if the array is empty.
-    bool empty() const { return Length == 0; }
-
-    const T *data() const { return Data; }
-
-    /// size - Get the array size.
-    size_t size() const { return Length; }
-
-    /// front - Get the first element.
-    const T &front() const {
-      AT_CHECK(!empty(), "ArrayRef: attempted to access front() of empty list");
-      return Data[0];
-    }
-
-    /// back - Get the last element.
-    const T &back() const {
-      AT_CHECK(!empty(), "ArrayRef: attempted to access back() of empty list");
-      return Data[Length-1];
-    }
-
-    /// equals - Check for element-wise equality.
-    bool equals(ArrayRef RHS) const {
-      if (Length != RHS.Length)
-        return false;
-      return std::equal(begin(), end(), RHS.begin());
-    }
-
-    /// slice(n, m) - Chop off the first N elements of the array, and keep M
-    /// elements in the array.
-    ArrayRef<T> slice(size_t N, size_t M) const {
-      AT_CHECK(N+M <= size(), "ArrayRef: invalid slice, ", N, " + ", M, " is not <= ", size());
-      return ArrayRef<T>(data()+N, M);
-    }
-
-    /// slice(n) - Chop off the first N elements of the array.
-    ArrayRef<T> slice(size_t N) const { return slice(N, size() - N); }
-
-    /// @}
-    /// @name Operator Overloads
-    /// @{
-    const T &operator[](size_t Index) const {
-      return Data[Index];
-    }
-
-    /// Vector compatibility
-    const T &at(size_t Index) const {
-      AT_CHECK(Index < Length, "ArrayRef: invalid index ", Index, " for length ", Length);
-      return Data[Index];
-    }
-
-    /// Disallow accidental assignment from a temporary.
-    ///
-    /// The declaration here is extra complicated so that "arrayRef = {}"
-    /// continues to select the move assignment operator.
-    template <typename U>
-    typename std::enable_if<std::is_same<U, T>::value, ArrayRef<T>>::type &
-    operator=(U &&Temporary) = delete;
-
-    /// Disallow accidental assignment from a temporary.
-    ///
-    /// The declaration here is extra complicated so that "arrayRef = {}"
-    /// continues to select the move assignment operator.
-    template <typename U>
-    typename std::enable_if<std::is_same<U, T>::value, ArrayRef<T>>::type &
-    operator=(std::initializer_list<U>) = delete;
-
-    /// @}
-    /// @name Expensive Operations
-    /// @{
-    std::vector<T> vec() const {
-      return std::vector<T>(Data, Data+Length);
-    }
-
-    /// @}
-    /// @name Conversion operators
-    /// @{
-    operator std::vector<T>() const {
-      return std::vector<T>(Data, Data+Length);
-    }
-
-    /// @}
-  };
-
-} // end namespace at
+#include <c10/util/ArrayRef.h>

aten/src/ATen/Backend.h

@@ -0,0 +1,2 @@
+#pragma once
+#include <c10/core/Backend.h>

aten/src/ATen/Backtrace.h

@@ -1,26 +1,2 @@
 #pragma once
-
-#include <cstddef>
-#include <string>
-#include <typeinfo>
-
-namespace at {
-/// Utility to demangle a C++ symbol name.
-std::string demangle(const char* name);
-
-/// Returns the printable name of the type.
-template <typename T>
-inline const char* demangle_type() {
-#ifdef __GXX_RTTI
-  static const std::string name = demangle(typeid(T).name());
-  return name.c_str();
-#else // __GXX_RTTI
-  return "(RTTI disabled, cannot show name)";
-#endif // __GXX_RTTI
-}
-
-std::string get_backtrace(
-    size_t frames_to_skip = 0,
-    size_t maximum_number_of_frames = 64,
-    bool skip_python_frames = true);
-} // namespace at
+#include <ATen/core/Backtrace.h>

aten/src/ATen/CMakeLists.txt

@@ -1,11 +1,6 @@
 cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
 SET(CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/cmake ${CMAKE_MODULE_PATH})
 
-if (NOT CAFFE2_CMAKE_BUILDING_WITH_MAIN_REPO)
-  # ---[ Generate and install header and cpp files
-  include(../../../cmake/Codegen.cmake)
-endif()
-
 IF(NOT MSVC)
   SET(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-ignored-qualifiers")
   SET(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -Wno-ignored-qualifiers")
@@ -13,23 +8,6 @@ IF(NOT MSVC)
   SET(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -Wno-absolute-value")
 ENDIF(NOT MSVC)
 
-################################################################################
-# Helper functions
-################################################################################
-
-function(filter_list output input)
-    unset(result)
-    foreach(filename ${${input}})
-        foreach(pattern ${ARGN})
-            if("${filename}" MATCHES "${pattern}")
-                list(APPEND result "${filename}")
-            endif()
-        endforeach()
-    endforeach()
-    set(${output} ${result} PARENT_SCOPE)
-endfunction()
-
-
 # Can be compiled standalone
 IF(NOT AT_INSTALL_BIN_DIR OR NOT AT_INSTALL_LIB_DIR OR NOT AT_INSTALL_INCLUDE_DIR OR NOT AT_INSTALL_SHARE_DIR)
   SET(AT_INSTALL_BIN_DIR "bin" CACHE PATH "AT install binary subdirectory")
@@ -42,13 +20,16 @@ CONFIGURE_FILE(Config.h.in "${CMAKE_CURRENT_SOURCE_DIR}/Config.h")
 CONFIGURE_FILE(cuda/CUDAConfig.h.in "${CMAKE_CURRENT_SOURCE_DIR}/cuda/CUDAConfig.h")
 
 # NB: If you edit these globs, you'll have to update setup.py package_data as well
-FILE(GLOB base_h "*.h" "detail/*.h")
-FILE(GLOB base_cpp "*.cpp" "detail/*.cpp")
+FILE(GLOB base_h "*.h" "detail/*.h" "cpu/*.h")
+FILE(GLOB base_cpp "*.cpp" "detail/*.cpp" "cpu/*.cpp")
+add_subdirectory(core)
 FILE(GLOB cuda_h "cuda/*.h" "cuda/detail/*.h" "cuda/*.cuh" "cuda/detail/*.cuh")
 FILE(GLOB cuda_cpp "cuda/*.cpp" "cuda/detail/*.cpp")
 FILE(GLOB cuda_cu "cuda/*.cu" "cuda/detail/*.cu")
 FILE(GLOB cudnn_h "cudnn/*.h" "cudnn/*.cuh")
 FILE(GLOB cudnn_cpp "cudnn/*.cpp")
+FILE(GLOB miopen_h "miopen/*.h")
+FILE(GLOB miopen_cpp "miopen/*.cpp")
 FILE(GLOB mkl_cpp "mkl/*.cpp")
 FILE(GLOB mkldnn_cpp "mkldnn/*.cpp")
 
@@ -57,12 +38,13 @@ FILE(GLOB native_sparse_cpp "native/sparse/*.cpp")
 FILE(GLOB native_sparse_cuda_cu "native/sparse/cuda/*.cu")
 FILE(GLOB native_sparse_cuda_cpp "native/sparse/cuda/*.cpp")
 FILE(GLOB native_cudnn_cpp "native/cudnn/*.cpp")
+FILE(GLOB native_miopen_cpp "native/miopen/*.cpp")
 FILE(GLOB native_cuda_cu "native/cuda/*.cu")
 FILE(GLOB native_cuda_cpp "native/cuda/*.cpp")
 FILE(GLOB native_mkl_cpp "native/mkl/*.cpp")
 FILE(GLOB native_mkldnn_cpp "native/mkldnn/*.cpp")
 
-set(all_cpu_cpp ${base_cpp} ${native_cpp} ${native_sparse_cpp} ${native_mkl_cpp} ${native_mkldnn_cpp} ${generated_cpp} ${ATen_CPU_SRCS} ${cpu_kernel_cpp})
+set(all_cpu_cpp ${base_cpp} ${ATen_CORE_SRCS} ${native_cpp} ${native_sparse_cpp} ${native_mkl_cpp} ${native_mkldnn_cpp} ${generated_cpp} ${ATen_CPU_SRCS} ${cpu_kernel_cpp})
 if(AT_MKL_ENABLED)
   set(all_cpu_cpp ${all_cpu_cpp} ${mkl_cpp})
 endif()
@@ -73,9 +55,14 @@ endif()
 IF(USE_CUDA OR USE_ROCM)
   list(APPEND ATen_CUDA_INCLUDE ${CMAKE_CURRENT_SOURCE_DIR}/cuda)
   set(ATen_CUDA_SRCS ${ATen_CUDA_SRCS} ${cuda_cu} ${native_cuda_cu} ${native_sparse_cuda_cu})
-  set(all_cuda_cpp ${native_cudnn_cpp} ${native_sparse_cuda_cpp} ${cuda_cpp} ${native_cuda_cpp} ${cuda_generated_cpp} ${ATen_CUDA_SRCS})
-  IF(CUDNN_FOUND)
-    SET(all_cuda_cpp ${all_cuda_cpp} ${cudnn_cpp})
+  set(all_cuda_cpp ${native_sparse_cuda_cpp} ${cuda_cpp} ${native_cuda_cpp} ${cuda_generated_cpp} ${ATen_CUDA_SRCS})
+  IF(USE_CUDA)
+    SET(all_cuda_cpp ${native_cudnn_cpp} ${native_miopen_cpp} ${all_cuda_cpp})
+    IF(CUDNN_FOUND)
+      SET(all_cuda_cpp ${all_cuda_cpp} ${cudnn_cpp})
+    ENDIF()
+  ELSEIF(USE_ROCM)
+    SET(all_cuda_cpp ${native_cudnn_cpp} ${native_miopen_cpp} ${miopen_cpp} ${all_cuda_cpp})
   ENDIF()
 endif()
 
@@ -161,7 +148,7 @@ endif(MKLDNN_FOUND)
 
 list(APPEND ATen_CPU_DEPENDENCY_LIBS cpuinfo)
 
-if(NOT MSVC)
+if(NOT MSVC AND NOT EMSCRIPTEN)
   # Preserve values for the main build
   set(__aten_sleef_build_shared_libs ${BUILD_SHARED_LIBS})
   set(__aten_sleef_build_tests ${BUILD_TESTS})
@@ -171,6 +158,16 @@ if(NOT MSVC)
   set(OLD_CMAKE_CXX_FLAGS ${CMAKE_CXX_FLAGS})
   set(CMAKE_CXX_FLAGS)
 
+  # Bump up optimization level for sleef to -O1, since at -O0 the compiler
+  # excessively spills intermediate vector registers to the stack
+  # and makes things run impossibly slowly
+  set(OLD_CMAKE_C_FLAGS_DEBUG ${CMAKE_C_FLAGS_DEBUG})
+  IF(${CMAKE_C_FLAGS_DEBUG} MATCHES "-O0")
+    string(REGEX REPLACE "-O0" "-O1" CMAKE_C_FLAGS_DEBUG ${OLD_CMAKE_C_FLAGS_DEBUG})
+  ELSE()
+    set(CMAKE_C_FLAGS_DEBUG "${CMAKE_C_FLAGS_DEBUG} -O1")
+  ENDIF()
+
   set(BUILD_SHARED_LIBS OFF CACHE BOOL "Build sleef static" FORCE)
   set(BUILD_DFT OFF CACHE BOOL "Don't build sleef DFT lib" FORCE)
   set(BUILD_GNUABI_LIBS OFF CACHE BOOL "Don't build sleef gnuabi libs" FORCE)
@@ -181,6 +178,7 @@ if(NOT MSVC)
   link_directories(${CMAKE_BINARY_DIR}/sleef/lib)
   list(APPEND ATen_CPU_DEPENDENCY_LIBS sleef)
 
+  set(CMAKE_C_FLAGS_DEBUG ${OLD_CMAKE_C_FLAGS_DEBUG})
   set(CMAKE_CXX_FLAGS ${OLD_CMAKE_CXX_FLAGS})
 
   # Set these back. TODO: Use SLEEF_ to pass these instead
@@ -208,6 +206,12 @@ IF(USE_CUDA AND NOT USE_ROCM)
 	--generate-code arch=compute_50,code=sm_50
 	--generate-code arch=compute_60,code=sm_60
 	--generate-code arch=compute_70,code=sm_70)
+    elseif(${CUDA_VERSION_MAJOR} EQUAL "10")
+      SET(CUFFT_FAKELINK_OPTIONS
+	--generate-code arch=compute_35,code=sm_35
+	--generate-code arch=compute_50,code=sm_50
+	--generate-code arch=compute_60,code=sm_60
+	--generate-code arch=compute_70,code=sm_70)
     else()
       MESSAGE(FATAL_ERROR "Unhandled major cuda version ${CUDA_VERSION_MAJOR}")
     endif()
@@ -250,15 +254,16 @@ IF(USE_CUDA AND NOT USE_ROCM)
   ENDIF(USE_MAGMA)
   IF ($ENV{ATEN_STATIC_CUDA})
     list(APPEND ATen_CUDA_DEPENDENCY_LIBS "${CUDA_TOOLKIT_ROOT_DIR}/lib64/libculibos.a")
+    list(APPEND ATen_CUDA_DEPENDENCY_LIBS "${CUDA_TOOLKIT_ROOT_DIR}/lib64/libcudart_static.a")
   ENDIF($ENV{ATEN_STATIC_CUDA})
 ENDIF()
 
 IF(USE_ROCM)
  ### Link in the ROCm libraries BLAS / RNG.
- FIND_LIBRARY(HIPBLAS_LIBRARY hipblas HINTS ${HIPBLAS_PATH}/lib)
- FIND_LIBRARY(HIPRNG_LIBRARY hcrng HINTS ${HIPRNG_PATH}/lib)
+ FIND_LIBRARY(ROCBLAS_LIBRARY rocblas HINTS ${ROCBLAS_PATH}/lib)
+ FIND_LIBRARY(HIPRAND_LIBRARY hiprand HINTS ${HIPRAND_PATH}/lib)
 
- list(APPEND ATen_CUDA_DEPENDENCY_LIBS ${HIPBLAS_LIBRARY} ${HIPRNG_LIBRARY})
+ list(APPEND ATen_CUDA_DEPENDENCY_LIBS ${ROCBLAS_LIBRARY} ${HIPRAND_LIBRARY})
 ENDIF()
 
 # Include CPU paths for CUDA as well
@@ -286,7 +291,7 @@ else()
   target_link_libraries(ATen_cpu PRIVATE ATEN_CPU_FILES_GEN_LIB)
   caffe2_interface_library(ATen_cpu ATen_cpu_library)
   # Set standard properties on the target
-  aten_set_target_props(ATen_cpu)
+  torch_set_target_props(ATen_cpu)
 
   # Make sure these don't get built by parent
   set(ATen_CPU_SRCS)
@@ -327,7 +332,7 @@ if(USE_CUDA OR USE_ROCM)
         ATen_cuda PUBLIC ATen_cpu ${ATen_PUBLIC_CUDA_DEPENDENCY_LIBS})
 
     # Set standard properties on the target
-    aten_set_target_props(ATen_cuda)
+    torch_set_target_props(ATen_cuda)
 
     caffe2_interface_library(ATen_cuda ATen_cuda_library)
 
@@ -345,9 +350,9 @@ if(NOT AT_LINK_STYLE STREQUAL "INTERFACE")
   endif()
 
   if(NOT MSVC)
-    aten_compile_options(ATen_cpu)
+    torch_compile_options(ATen_cpu)
     if(USE_CUDA OR USE_ROCM)
-      aten_compile_options(ATen_cuda)
+      torch_compile_options(ATen_cuda)
     endif()
   endif()
 
@@ -359,41 +364,13 @@ if(NOT AT_LINK_STYLE STREQUAL "INTERFACE")
   endif()
 endif()
 
-if (NOT CAFFE2_CMAKE_BUILDING_WITH_MAIN_REPO)
-  # Eventually replace this use of LOCATION with use of
-  # $<TARGET_FILE:ATen_cpu>, but generators only work in some cases
-  cmake_policy(SET CMP0026 OLD)
-  get_target_property(ATEN_CPU_OUTPUT_NAME ATen_cpu LOCATION)
-  get_filename_component(ATEN_CPU_OUTPUT_NAME ${ATEN_CPU_OUTPUT_NAME} NAME)
-  set(ATEN_LIBRARIES
-    "${CMAKE_INSTALL_PREFIX}/${AT_INSTALL_LIB_DIR}/${ATEN_CPU_OUTPUT_NAME}")
-  if(USE_CUDA OR USE_ROCM)
-    get_target_property(ATEN_CUDA_OUTPUT_NAME ATen_cuda LOCATION)
-    get_filename_component(ATEN_CUDA_OUTPUT_NAME ${ATEN_CUDA_OUTPUT_NAME} NAME)
-    list(APPEND ATEN_LIBRARIES
-      "${CMAKE_INSTALL_PREFIX}/${AT_INSTALL_LIB_DIR}/${ATEN_CUDA_OUTPUT_NAME}")
-  endif()
-
-  install(TARGETS ATen_cpu
-    RUNTIME DESTINATION "${AT_INSTALL_BIN_DIR}"
-    LIBRARY DESTINATION "${AT_INSTALL_LIB_DIR}"
-    ARCHIVE DESTINATION "${AT_INSTALL_LIB_DIR}")
-
-  if(USE_CUDA OR USE_ROCM)
-    install(TARGETS ATen_cuda
-      RUNTIME DESTINATION "${AT_INSTALL_BIN_DIR}"
-      LIBRARY DESTINATION "${AT_INSTALL_LIB_DIR}"
-      ARCHIVE DESTINATION "${AT_INSTALL_LIB_DIR}")
-  endif()
-endif()
-
 SET(ATEN_INCLUDE_DIR "${CMAKE_INSTALL_PREFIX}/${AT_INSTALL_INCLUDE_DIR}")
 CONFIGURE_FILE(ATenConfig.cmake.in "${CMAKE_CURRENT_BINARY_DIR}/cmake-exports/ATenConfig.cmake")
 INSTALL(FILES "${CMAKE_CURRENT_BINARY_DIR}/cmake-exports/ATenConfig.cmake"
   DESTINATION "${AT_INSTALL_SHARE_DIR}/cmake/ATen")
 
 # https://stackoverflow.com/questions/11096471/how-can-i-install-a-hierarchy-of-files-using-cmake
-FOREACH(HEADER ${base_h} ${cuda_h} ${cudnn_h})
+FOREACH(HEADER ${base_h} ${ATen_CORE_HEADERS} ${cuda_h} ${cudnn_h})
   string(REPLACE "${CMAKE_CURRENT_SOURCE_DIR}/" "" HEADER_SUB ${HEADER})
   GET_FILENAME_COMPONENT(DIR ${HEADER_SUB} DIRECTORY)
   INSTALL(FILES ${HEADER} DESTINATION ${AT_INSTALL_INCLUDE_DIR}/ATen/${DIR})
@@ -411,39 +388,8 @@ else()
   add_subdirectory(test)
 endif()
 
-if (NOT CAFFE2_CMAKE_BUILDING_WITH_MAIN_REPO)
-  foreach(test_src ${ATen_CPU_TEST_SRCS})
-    get_filename_component(test_name ${test_src} NAME_WE)
-    add_executable(${test_name} "${test_src}")
-    target_include_directories(
-        ${test_name} PRIVATE $<INSTALL_INTERFACE:include>)
-    target_include_directories(${test_name} PRIVATE ${ATen_CPU_INCLUDE})
-    target_include_directories(${test_name} SYSTEM PRIVATE ${ATen_THIRD_PARTY_INCLUDE})
-    target_link_libraries(${test_name} ATen_cpu)
-    add_test(NAME ${test_name} COMMAND $<TARGET_FILE:${test_name}>)
-    install(TARGETS ${test_name} DESTINATION test)
-  endforeach()
-
-  if(USE_CUDA OR USE_ROCM)
-    foreach(test_src ${ATen_CUDA_TEST_SRCS})
-      get_filename_component(test_name ${test_src} NAME_WE)
-      torch_cuda_based_add_executable(${test_name} "${test_src}")
-      target_include_directories(
-          ${test_name} PRIVATE $<INSTALL_INTERFACE:include>)
-      target_include_directories(${test_name} PRIVATE ${ATen_CPU_INCLUDE})
-      target_include_directories(${test_name} SYSTEM PRIVATE ${ATen_THIRD_PARTY_INCLUDE})
-      target_link_libraries(${test_name} -Wl,--no-as-needed ATen_cpu ATen_cuda)
-      add_test(NAME ${test_name} COMMAND $<TARGET_FILE:${test_name}>)
-      install(TARGETS ${test_name} DESTINATION test)
-    endforeach()
-  endif()
-
-  # Make sure these don't get built by parent
-  set(ATen_CPU_TEST_SRCS)
-  set(ATen_CUDA_TEST_SRCS)
-endif()
-
 # Pass source, includes, and libs to parent
+set(ATen_CORE_SRCS ${ATen_CORE_SRCS} PARENT_SCOPE)
 set(ATen_CPU_SRCS ${ATen_CPU_SRCS} PARENT_SCOPE)
 set(ATen_CUDA_SRCS ${ATen_CUDA_SRCS} PARENT_SCOPE)
 set(ATen_CPU_TEST_SRCS ${ATen_CPU_TEST_SRCS} PARENT_SCOPE)

aten/src/ATen/CPUApplyUtils.h

@@ -3,9 +3,99 @@
 #include "ATen/Parallel.h"
 #include "ATen/TensorUtils.h"
 #include <limits>
+#include <utility>
+#include <cstring>
 
 namespace at {
 
+/*
+[collapse dims] Updates sizes, and strides to reflect a "collapse" of
+the info, possibly excluding the optional excludeDim. A "collapsed" version
+of the info is the fewest dims that order the tensor's elements in the same
+way as the original info. If excludeDim is specified, the collapse is the
+fewest dims that order the tensor's elements as the original and preserve the
+excluded dimension, unless the tensor collapses to a point.
+
+This function returns a pair of values.
+
+1) The (new) index of the preserved dimension if excludeDim is
+specified. 0 if the tensor is collapsed to a point. -1
+otherwise.
+
+2) The new number of dimensions.
+*/
+template <typename T>
+inline std::pair<int64_t, int64_t> collapse_dims(
+    T* sizes,
+    T* strides,
+    int64_t dims,
+    const int excludeDim = -1) {
+  AT_CHECK(
+      excludeDim >= -1 && excludeDim < dims,
+      "expected excluded dim between -1 and dims - 1");
+
+  int64_t stopDim = (excludeDim == -1) ? dims : excludeDim;
+  int64_t newIndex = -1;
+  int64_t oldIndex = 0;
+  int64_t remappedExcludedDim = -1;
+
+  while (oldIndex < dims) {
+    // Finds a dimension to collapse into
+    for (; oldIndex < stopDim; ++oldIndex) {
+      if (sizes[oldIndex] == 1) {
+        continue;
+      }
+
+      ++newIndex;
+      sizes[newIndex] = sizes[oldIndex];
+      strides[newIndex] = strides[oldIndex];
+      ++oldIndex;
+      break;
+    }
+
+    // Collapses dims
+    for (; oldIndex < stopDim; ++oldIndex) {
+      if (sizes[oldIndex] == 1) {
+        continue;
+      }
+
+      if (strides[newIndex] == sizes[oldIndex] * strides[oldIndex]) {
+        sizes[newIndex] *= sizes[oldIndex];
+        strides[newIndex] = strides[oldIndex];
+      } else {
+        ++newIndex;
+        sizes[newIndex] = sizes[oldIndex];
+        strides[newIndex] = strides[oldIndex];
+      }
+    }
+
+    // Handles excludeDim being set (oldIndex == excludeDim)
+    if (oldIndex != dims) {
+      // Preserves excluded dimension
+      ++newIndex;
+      sizes[newIndex] = sizes[oldIndex];
+      strides[newIndex] = strides[oldIndex];
+      remappedExcludedDim = newIndex;
+
+      // Restarts iteration after excludeDim
+      ++oldIndex;
+      stopDim = dims;
+    }
+  }
+
+  // Handles special case of all dims size 1
+  if (newIndex == -1 || (newIndex == 0 && sizes[0] == 1)) {
+    dims = 1;
+    sizes[0] = 1;
+    strides[0] = 1;
+
+    return std::pair<int64_t, int64_t>(0, 1);
+  }
+
+  dims = newIndex + 1;
+  return std::pair<int64_t, int64_t>(remappedExcludedDim, dims);
+}
+
 /*
  * The basic strategy for apply is as follows:
  *
@@ -50,27 +140,6 @@ inline Tensor sort_strides(Tensor& tensor_) {
   return tensor;
 }
 
-template <typename Arg>
-inline void _setup_arrays(Tensor& tensor, Arg* iter) {
-  int64_t max_dim = tensor.ndimension();
-  iter->dim_ = 0;
-  for (int64_t i = 0; i < max_dim; i++) {
-    int64_t size = tensor.size(i);
-    int64_t stride = tensor.stride(i);
-    while (i + 1 < max_dim &&
-           (tensor.size(i + 1) == 1 ||
-            tensor.stride(i) == tensor.size(i + 1) * tensor.stride(i + 1))) {
-      size = size * tensor.size(i + 1);
-      if (tensor.size(i + 1) != 1)
-        stride = tensor.stride(i + 1);
-      i++;
-    }
-    iter->sizes_[iter->dim_] = size;
-    iter->strides_[iter->dim_] = stride;
-    iter->dim_++;
-  }
-}
-
 template <typename T, int N>
 struct strided_tensor_iter_fixed {
  public:
@@ -86,8 +155,16 @@ struct strided_tensor_iter_fixed {
   strided_tensor_iter_fixed(strided_tensor_iter_fixed&&) = default;
   strided_tensor_iter_fixed(Tensor& tensor, bool sort_strides = false)
       : data_(tensor.data<T>()) {
-    memset(counter_, 0, sizeof(int64_t) * N);
-    _setup_arrays(tensor, this);
+    std::memset(counter_, 0, sizeof(int64_t) * N);
+    if (tensor.dim() > 0) {
+      std::memcpy(
+          sizes_, tensor.sizes().data(), tensor.dim() * sizeof(int64_t));
+      std::memcpy(
+          strides_,
+          tensor.strides().data(),
+          tensor.dim() * sizeof(int64_t));
+    }
+    dim_ = std::get<1>(collapse_dims(sizes_, strides_, tensor.ndimension()));
   }
 };
 
@@ -109,9 +186,9 @@ struct strided_tensor_iter {
       : data_(tensor.data<T>()),
         dim_(tensor.ndimension()),
         counter_(dim_, 0),
-        sizes_(tensor.sizes()),
-        strides_(tensor.strides()) {
-    _setup_arrays(tensor, this);
+        sizes_(tensor.sizes().vec()),
+        strides_(tensor.strides().vec()) {
+    dim_ = std::get<1>(collapse_dims(sizes_.data(), strides_.data(), dim_));
   }
 };
 
@@ -132,7 +209,7 @@ inline std::string _all_equal_numel_error(at::ArrayRef<Tensor> tensors) {
   for (size_t i = 0; i < tensors.size() - 1; i++) {
     oss << tensors[i].sizes() << ", ";
   }
-  oss << "and " << tensors[tensors.size() - 1]
+  oss << "and " << tensors[tensors.size() - 1].sizes()
       << " to have the same number of elements, but got ";
   for (size_t i = 0; i < tensors.size() - 1; i++) {
     oss << tensors[i].numel() << ", ";
@@ -145,7 +222,7 @@ inline std::string _all_equal_numel_error(at::ArrayRef<Tensor> tensors) {
 inline bool _apply_preamble(ArrayRef<Tensor> tensors) {
   checkBackend("CPU_tensor_apply", tensors, Backend::CPU);
   if (!_all_equal_numel(tensors))
-    throw std::runtime_error(_all_equal_numel_error(tensors));
+    AT_ERROR(_all_equal_numel_error(tensors));
   // An empty tensor has no elements
   for (auto& t : tensors)
     if (t.numel() == 0)

aten/src/ATen/CPUFixedAllocator.h

@@ -1,7 +1,7 @@
 #pragma once
 
 #include "TH/TH.h"
-#include "ATen/Error.h"
+#include "c10/util/Exception.h"
 
 // This file creates a fake allocator that just throws exceptions if
 // it is actually used.

aten/src/ATen/CPUGeneral.h

@@ -1,12 +1,12 @@
 #pragma once
 
-// Using AT_API is crucial as otherwise you'll see
+// Using CAFFE2_API is crucial as otherwise you'll see
 // linking errors using MSVC
 // See https://msdn.microsoft.com/en-us/library/a90k134d.aspx
-// This header adds this if using AT_API
-#include "ATen/ATenGeneral.h"
+// This header adds this if using CAFFE2_API
+#include "ATen/core/ATenGeneral.h"
 
 namespace at {
-AT_API void set_num_threads(int);
-AT_API int get_num_threads();
+CAFFE2_API void set_num_threads(int);
+CAFFE2_API int get_num_threads();
 }

aten/src/ATen/CPUTypeDefault.cpp

@@ -0,0 +1,20 @@
+#include <ATen/CPUTypeDefault.h>
+
+#include <ATen/Context.h>
+#include <ATen/CPUGenerator.h>
+
+namespace at {
+
+Allocator* CPUTypeDefault::allocator() const {
+  return getCPUAllocator();
+}
+
+Device CPUTypeDefault::getDeviceFromPtr(void * data) const {
+  return DeviceType::CPU;
+}
+
+std::unique_ptr<Generator> CPUTypeDefault::generator() const {
+  return std::unique_ptr<Generator>(new CPUGenerator(&at::globalContext()));
+}
+
+} // namespace at

aten/src/ATen/CPUTypeDefault.h

@@ -0,0 +1,14 @@
+#pragma once
+#include <ATen/TypeDefault.h>
+
+namespace at {
+
+struct CAFFE2_API CPUTypeDefault : public TypeDefault {
+  CPUTypeDefault(TensorTypeId type_id, bool is_variable, bool is_undefined)
+      : TypeDefault(type_id, is_variable, is_undefined) {}
+  Allocator* allocator() const override;
+  Device getDeviceFromPtr(void * data) const override;
+  std::unique_ptr<Generator> generator() const override;
+};
+
+} // namespace at

aten/src/ATen/CUDAStream.cpp

@@ -1,183 +0,0 @@
-#include "ATen/CUDAStream.h"
-#include "ATen/Error.h"
-#include "ATen/detail/CUDAHooksInterface.h"
-
-#include <mutex>
-
-// Internal implementation is entirely hidden 
-struct CUDAStreamInternals {
-  bool is_destructible;
-  std::atomic<int> refcount;
-  int64_t device; // Note: cudaGetDevice works with int32_t, not int64_t
-  cudaStream_t stream;
-};
-
-namespace at {
-
-namespace detail {
-
-  /*
-  * Stream state
-  */
-  static constexpr cudaStream_t DEFAULT_STREAM = 0;
-
-  static std::once_flag init_flag;
-  static int64_t num_gpus;
-  static CUDAStreamInternals* default_streams;
-  static thread_local CUDAStreamInternals** current_streams = nullptr;
-
-  // Creates a(n indestructible) default stream for each device
-  // Note: the default stream on each device is signified by a zero
-  // value for the pointer, and so is not actually created as usual.
-  // In particular, we don't need to switch devices when creating the 
-  // streams.
-  static void initDefaultCUDAStreams() {
-    num_gpus = getCUDAHooks().getNumGPUs();
-    default_streams = (CUDAStreamInternals*) malloc(num_gpus * sizeof(CUDAStreamInternals));
-    for (auto i = decltype(num_gpus){0}; i < num_gpus; ++i) {
-      default_streams[i].is_destructible = false;
-      default_streams[i].refcount = 0;
-      default_streams[i].device = i;
-      default_streams[i].stream = DEFAULT_STREAM;
-    }
-  }
-
-  // Init front-end to ensure initialization only occurs once
-  static void initCUDAStreamsOnce() {
-    // Inits default streams (once, globally)
-    std::call_once(init_flag, initDefaultCUDAStreams);
-    
-    // Inits current streams (thread local) to default streams    
-    if (current_streams) return;
-    current_streams = (CUDAStreamInternals**) malloc(num_gpus * sizeof(CUDAStreamInternals*));
-    for (auto i = decltype(num_gpus){0}; i < num_gpus; ++i) {
-      current_streams[i] = &default_streams[i];
-    }
-  }
-
-  /*
-  * Pointer-based stream API
-  */
-
-  // Helper to return the current device
-  static inline int64_t current_device() {
-    int cur_device;
-    DynamicCUDAInterface::get_device(&cur_device);
-    return cur_device;
-  }
-
-  // Helper to verify the GPU index is valid
-  static inline void check_gpu(int64_t device) {
-    AT_CHECK(device >= 0 && device < num_gpus);
-  }
-
-  CUDAStreamInternals* CUDAStream_getDefaultStreamOnDevice(int64_t device) {
-    initCUDAStreamsOnce();
-    check_gpu(device);
-    return &default_streams[device];
-  }
-  CUDAStreamInternals* CUDAStream_getDefaultStream() {
-    return CUDAStream_getDefaultStreamOnDevice(current_device());
-  }
-
-  // Creates (and retains) and new cuda stream
-  CUDAStreamInternals* CUDAStream_createAndRetainWithOptions(int32_t flags, int32_t priority) {
-    CUDAStreamInternals* internals = (CUDAStreamInternals*) malloc(sizeof(CUDAStreamInternals));
-    internals->is_destructible = true;
-    internals->refcount = 1;
-    internals->device = current_device();
-    DynamicCUDAInterface::cuda_stream_create_with_priority(&internals->stream, flags, priority);
-    return internals;
-  }
-
-  // Note: despite not being "unsafe," is using these methods in a multithreaded
-  // environment then the caller must be sure that streams are valid 
-  // when they're requested. These methods will throw an error if an
-  // invalid stream is requested.
-  CUDAStreamInternals* CUDAStream_getAndRetainCurrentStreamOnDevice(int64_t device) {
-    initCUDAStreamsOnce();
-    check_gpu(device);
-    auto cur = current_streams[device];
-    AT_CHECK(CUDAStream_retain(cur));
-    return cur;
-  }
-  CUDAStreamInternals* CUDAStream_getAndRetainCurrentStream() {
-    return CUDAStream_getAndRetainCurrentStreamOnDevice(current_device());
-  }
-
-  // Note: these unsafe methods do not retain the stream before returning it.
-  // This is unsafe behavior and these methods SHOULD NOT BE USED.
-  // They are here only for legacy compatibility.
-  CUDAStreamInternals* CUDAStream_getCurrentStreamOnDeviceUnsafe(int64_t device) {
-    initCUDAStreamsOnce();
-    check_gpu(device);
-    return current_streams[device];
-  }
-  CUDAStreamInternals* CUDAStream_getCurrentStreamUnsafe() {
-    return CUDAStream_getCurrentStreamOnDeviceUnsafe(current_device());
-  }
-
-  void CUDAStream_setStreamOnDevice(int64_t device, CUDAStreamInternals* ptr) {
-    initCUDAStreamsOnce();
-    check_gpu(device);
-    AT_CHECK(ptr);
-    AT_CHECK(ptr->device == device);
-    AT_CHECK(CUDAStream_retain(ptr));
-
-    CUDAStream_free(current_streams[device]);
-    current_streams[device] = ptr;
-  }
-  void CUDAStream_setStream(CUDAStreamInternals* ptr) {
-    CUDAStream_setStreamOnDevice(current_device(), ptr);
-  }
-
-  // Getters
-  cudaStream_t CUDAStream_stream(CUDAStreamInternals* ptr) {
-    AT_CHECK(ptr);
-    return ptr->stream;
-  }
-
-  int64_t CUDAStream_device(CUDAStreamInternals* ptr) {
-    AT_CHECK(ptr);
-    return ptr->device;
-  }
-
-  // Memory management
-  // Note: only destructible (non-default) streams are ref counted
-  bool CUDAStream_retain(CUDAStreamInternals* ptr) {
-    AT_CHECK(ptr);
-    if (ptr->is_destructible) return(++ptr->refcount > 1);
-    return true;
-  }
-
-  void CUDAStream_free(CUDAStreamInternals*& ptr) {
-    if (ptr && ptr->stream && ptr->is_destructible && --ptr->refcount <= 0) {
-      AT_CHECK(ptr->refcount == 0);
-      DynamicCUDAInterface::cuda_stream_destroy(ptr->stream);
-      free(ptr);
-      ptr = nullptr;
-    }
-  }
-
-} // namespace detail
-
-  /*
-  * CUDAStream functions
-  */
-
-   // Copy constructor
-  CUDAStream::CUDAStream(const CUDAStream& other) {
-    AT_CHECK(other.internals_);
-    AT_CHECK(detail::CUDAStream_retain(other.internals_));
-
-    internals_ = other.internals_;
-  }
-
-  // Move constructor
-  CUDAStream::CUDAStream(CUDAStream&& other) {
-    AT_CHECK(other.internals_);
-
-    std::swap(internals_, other.internals_);
-  }
-  
-} // namespace at

aten/src/ATen/CUDAStream.h

@@ -1,95 +0,0 @@
-#pragma once
-
-#include <cstdint>
-#include <utility>
-
-/*
-* A CUDA stream interface with no CUDA build dependency.
-* 
-* Includes the CUDAStream RAII class and a pointer-based stream API.
-* 
-* The ATen Context interface should be preferred when working with streams.
-*/
-
-// Forward-declares cudaStream_t to avoid depending on CUDA in CPU builds
-// Note: this is the internal CUDA runtime typedef for cudaStream_t
-struct CUstream_st;
-typedef struct CUstream_st* cudaStream_t;
-
-// Forward-declares internals
-struct CUDAStreamInternals;
-
-namespace at {
-
-namespace detail {
-
-// Pointer-based API (for internal use)
-// Note: ATen/Context is preferred to work with streams safely
-CUDAStreamInternals* CUDAStream_getDefaultStreamOnDevice(int64_t device);
-CUDAStreamInternals* CUDAStream_getDefaultStream();
-
-CUDAStreamInternals* CUDAStream_createAndRetainWithOptions(int32_t flags, int32_t priority);
-
-CUDAStreamInternals* CUDAStream_getAndRetainCurrentStreamOnDevice(int64_t device);
-CUDAStreamInternals* CUDAStream_getAndRetainCurrentStream();
-
-// Note: these Unsafe gets should NEVER be used and are only here for legacy
-// purposes. Once those uses are gone they should be removed.
-CUDAStreamInternals* CUDAStream_getCurrentStreamOnDeviceUnsafe(int64_t device);
-CUDAStreamInternals* CUDAStream_getCurrentStreamUnsafe();
-
-void CUDAStream_setStreamOnDevice(int64_t device, CUDAStreamInternals* internals);
-void CUDAStream_setStream(CUDAStreamInternals* internals);
-
-cudaStream_t CUDAStream_stream(CUDAStreamInternals*);
-int64_t CUDAStream_device(CUDAStreamInternals*);
-
-bool CUDAStream_retain(CUDAStreamInternals*);
-void CUDAStream_free(CUDAStreamInternals*&);
-
-} // namespace detail
-
-// RAII for a CUDA stream
-// Allows use as a cudaStream_t, copying, moving, and metadata access.
-struct CUDAStream {
-  // Constants
-  static constexpr int32_t DEFAULT_FLAGS = 1; // = cudaStreamNonBlocking;
-  static constexpr int32_t DEFAULT_PRIORITY = 0;
-
-  // Constructors
-  CUDAStream() = default;
-  CUDAStream(CUDAStreamInternals* internals) : internals_{internals} { }
-  
-  // Destructor
-  ~CUDAStream() { detail::CUDAStream_free(internals_); }
-
-  // Copy constructor
-  CUDAStream(const CUDAStream& other);
-
-  // Move constructor
-  CUDAStream(CUDAStream&& other);  
-
-  // Assignment operator
-  CUDAStream& operator=(CUDAStream other) {
-    std::swap(internals_, other.internals_);
-    return *this;
-  }
-
-  // Implicit conversion to cudaStream_t
-  operator cudaStream_t() const { return detail::CUDAStream_stream(internals_); }
-
-  // Less than operator (to allow use in sets)
-  friend bool operator<(const CUDAStream& left, const CUDAStream& right) {
-    return left.internals_ < right.internals_;
-  }
-
-  // Getters
-  int64_t device() const { return detail::CUDAStream_device(internals_); }
-  cudaStream_t stream() const { return detail::CUDAStream_stream(internals_); }
-  CUDAStreamInternals* internals() const { return internals_; }
-
-private:
-  CUDAStreamInternals* internals_ = nullptr;
-};
-
-} // namespace at

aten/src/ATen/CheckGenerator.h

@@ -1,8 +1,8 @@
 #pragma once
 
-#include "ATen/Error.h"
-#include "ATen/Generator.h"
 #include "ATen/Utils.h"
+#include "ATen/core/Generator.h"
+#include "c10/util/Exception.h"
 
 namespace at {
 

aten/src/ATen/Config.h.in

@@ -8,3 +8,4 @@
 
 #define AT_MKLDNN_ENABLED() @AT_MKLDNN_ENABLED@
 #define AT_MKL_ENABLED() @AT_MKL_ENABLED@
+#define CAFFE2_STATIC_LINK_CUDA() @CAFFE2_STATIC_LINK_CUDA@

aten/src/ATen/Context.cpp

@@ -2,6 +2,8 @@
 
 #include "Context.h"
 
+#include <ATen/core/TensorOptions.h>
+
 #include <thread>
 #include <mutex>
 #include <sstream>
@@ -9,10 +11,11 @@
 #include <stdexcept>
 
 #include "ATen/CPUGenerator.h"
+#include "ATen/RegisterCPU.h"
+#include "ATen/Tensor.h"
+#include <ATen/cpu/FlushDenormal.h>
 
-#ifdef USE_SSE3
-#include <pmmintrin.h>
-#endif
+#include "TH/TH.h"  // for USE_LAPACK
 
 namespace at {
 
@@ -27,16 +30,20 @@ static inline void argErrorHandler(int arg, const char * msg, void * data) {
 
 Context::Context()
 : next_id(static_cast<size_t>(TypeID::NumOptions))
-, thc_state(nullptr, [](THCState* p){ /* no-op */ } ) {
+, thc_state(nullptr, [](THCState* p){ /* no-op */ } )
+, thh_state(nullptr, [](THHState* p){ /* no-op */ } )
+{
 
   THSetDefaultErrorHandler(errorHandler,nullptr);
   THSetDefaultArgErrorHandler(argErrorHandler,nullptr);
 
-  generator_registry[static_cast<int>(Backend::CPU)]
+  generator_registry[static_cast<int>(DeviceType::CPU)]
     .reset(new CPUGenerator(this));
-  Type::registerCPU(this);
+  register_cpu_types(this);
 }
 
+// TODO: This could be bad juju if someone calls globalContext() in the
+// destructor of an object with static lifetime.
 Context & globalContext() {
   static Context globalContext_;
   return globalContext_;
@@ -77,19 +84,63 @@ bool Context::hasMKL() const {
 #endif
 }
 
-bool Context::setFlushDenormal(bool on) {
-#ifdef USE_SSE3
-  // Setting flush-to-zero (FTZ) flag
-  _MM_SET_FLUSH_ZERO_MODE(on ? _MM_FLUSH_ZERO_ON
-                             : _MM_FLUSH_ZERO_OFF);
-
-  // Setting denormals-are-zero (DAZ) flag
-  _MM_SET_DENORMALS_ZERO_MODE(on ? _MM_DENORMALS_ZERO_ON
-                                 : _MM_DENORMALS_ZERO_OFF);
+bool Context::hasLAPACK() const {
+#ifdef USE_LAPACK
   return true;
 #else
   return false;
 #endif
 }
 
+bool Context::setFlushDenormal(bool on) {
+  return at::cpu::set_flush_denormal(on);
+}
+
+TypeExtendedInterface& getType(TensorOptions options) {
+  return globalContext().getType(
+            options.backend(), typeMetaToScalarType(options.dtype()), options.is_variable());
+}
+
+TypeExtendedInterface& getType(const TensorImpl* impl) {
+  Backend backend = tensorTypeIdToBackend(impl->type_id());
+  return globalContext().getType(
+            backend, typeMetaToScalarType(impl->dtype()), impl->is_variable());
+}
+
+TypeExtendedInterface& getType(const Tensor& t) {
+  return getType(t.unsafeGetTensorImpl());
+}
+
+LegacyTHDispatcher& getLegacyTHDispatcher(TensorOptions options) {
+  return globalContext().getLegacyTHDispatcher(
+            options.backend(), typeMetaToScalarType(options.dtype()));
+}
+
+LegacyTHDispatcher& getLegacyTHDispatcher(const TensorImpl* impl) {
+  Backend backend = tensorTypeIdToBackend(impl->type_id());
+  return globalContext().getLegacyTHDispatcher(
+            backend, typeMetaToScalarType(impl->dtype()));
+}
+
+Allocator* getCPUAllocator() {
+  return getTHDefaultAllocator();
+}
+
+struct LegacyDeviceTypeInit : public LegacyDeviceTypeInitInterface {
+  LegacyDeviceTypeInit(LegacyDeviceTypeInitArgs) {}
+  void initCPU() const override {
+    globalContext();
+  }
+  void initCUDA() const override {
+    globalContext().lazyInitCUDA();
+  }
+  void initHIP() const override {
+    globalContext().lazyInitHIP();
+  }
+  void initComplex() const override {
+    globalContext().lazyInitComplex();
+  }
+};
+REGISTER_LEGACY_TYPE_INIT(LegacyDeviceTypeInit);
+
 }

aten/src/ATen/Context.h

@@ -1,13 +1,19 @@
 #pragma once
 
-#include "ATen/ATenGeneral.h"
 #include <ATen/CPUGeneral.h>
-#include "ATen/Generator.h"
 #include "ATen/Type.h"
+#include "ATen/TypeExtendedInterface.h"
 #include "ATen/Utils.h"
-#include "ATen/Error.h"
+#include "ATen/LegacyTHDispatch.h"
+#include "ATen/LegacyTHDispatcher.h"
+#include "ATen/core/ATenGeneral.h"
+#include "ATen/core/Generator.h"
+#include "ATen/core/LegacyTypeDispatch.h"
+#include "ATen/core/VariableHooksInterface.h"
 #include "ATen/detail/CUDAHooksInterface.h"
-#include "ATen/CUDAStream.h"
+#include "ATen/detail/HIPHooksInterface.h"
+#include "ATen/detail/ComplexHooksInterface.h"
+#include "c10/util/Exception.h"
 
 #include <memory>
 #include <mutex>
@@ -15,119 +21,94 @@
 
 namespace at {
 
-enum class IsVariable {
-  NotVariable,
-  Variable,
-  NumOptions
-};
+class Tensor;
 
-class AT_API Context {
-public:
+class CAFFE2_API Context {
+ public:
   Context();
-  Type* getTypeRaw(Backend p, ScalarType s) {
-    return type_registry[static_cast<int>(p)][static_cast<int>(s)].get();
+  TypeExtendedInterface* getNonVariableTypeRaw(Backend p, ScalarType s) {
+    return static_cast<TypeExtendedInterface*>(globalLegacyTypeDispatch().getNonVariableTypeRaw(p, s));
+  }
+  TypeExtendedInterface * getNonVariableTypeOpt(Backend p, ScalarType s) {
+    return static_cast<TypeExtendedInterface*>(globalLegacyTypeDispatch().getNonVariableTypeOpt(p, s));
+  }
+  TypeExtendedInterface & getNonVariableType(Backend p, ScalarType s) {
+    return static_cast<TypeExtendedInterface&>(globalLegacyTypeDispatch().getNonVariableType(p, s));
+  }
+  TypeExtendedInterface & getVariableType(Backend p, ScalarType s) {
+    return static_cast<TypeExtendedInterface&>(globalLegacyTypeDispatch().getVariableType(p, s));
+  }
+  TypeExtendedInterface & getType(Backend p, ScalarType s, bool is_variable) {
+    return static_cast<TypeExtendedInterface&>(globalLegacyTypeDispatch().getType(p, s, is_variable));
+  }
+  LegacyTHDispatcher& getLegacyTHDispatcher(Backend p, ScalarType s) {
+    return globalLegacyTHDispatch().getLegacyTHDispatcher(p, s);
+  }
+  // The passed in Type must be delete'able
+  // TODO: Just make it take a unique_ptr
+  void registerType(Backend b, ScalarType s, Type* t) {
+    globalLegacyTypeDispatch().registerType(b, s,
+      LegacyTypeDispatch::TypeUniquePtr{t, LegacyTypeDeleter([](Type* p) { delete p; }) });
   }
-  Type * getTypeOpt(Backend p, ScalarType s) {
-    initCUDAIfNeeded(p);
-    auto type = getTypeRaw(p, s);
 
-    if(!type) {
-      // there is only a single Undefined Type.
-      if (p == Backend::Undefined || s == ScalarType::Undefined) {
-        return getTypeRaw(Backend::Undefined, ScalarType::Undefined);
-      }
-    }
+  void registerLegacyTHDispatcher(Backend b, ScalarType s, LegacyTHDispatcher* t) {
+    globalLegacyTHDispatch().registerDispatcher(b, s,
+      LegacyTHDispatch::LegacyTHDispatcherUniquePtr{t, LegacyTHDispatcherDeleter([](LegacyTHDispatcher* p) { delete p; }) });
+  }
 
-    return type;
-  }
-  Type & getType(Backend p, ScalarType s) {
-    auto* type = getTypeOpt(p, s);
-    if (!type) AT_ERROR(toString(p), toString(s), "Type is not enabled.");
-    return *type;
-  }
-  Generator & defaultGenerator(Backend p) {
-    initCUDAIfNeeded(p);
-    auto & generator = generator_registry[static_cast<int>(p)];
+  Generator & defaultGenerator(DeviceType device_type) {
+    initCUDAIfNeeded(device_type);
+    initHIPIfNeeded(device_type);
+    auto & generator = generator_registry[static_cast<int>(device_type)];
     if(!generator)
-      AT_ERROR(toString(p), " backend type not enabled.");
+      AT_ERROR(DeviceTypeName(device_type), " backend type not enabled.");
     return *generator;
   }
   bool hasMKL() const;
+  bool hasLAPACK() const;
+  bool hasMAGMA() const {
+    return detail::getCUDAHooks().hasMAGMA();
+  }
   bool hasCUDA() const {
     return detail::getCUDAHooks().hasCUDA();
   }
-  bool hasCuDNN() const {
-    return detail::getCUDAHooks().hasCuDNN();
+  bool hasHIP() const {
+    return detail::getHIPHooks().hasHIP();
   }
-  int64_t current_device() const {
-    return detail::getCUDAHooks().current_device();
-  }
-  // defined in header so that getType has ability to inline
-  // call_once check. getType is called fairly frequently
+  // defined in header so that getNonVariableType has ability to inline
+  // call_once check. getNonVariableType is called fairly frequently
   THCState* lazyInitCUDA() {
     std::call_once(thc_init,[&] {
       thc_state = detail::getCUDAHooks().initCUDA();
-      generator_registry[static_cast<int>(Backend::CUDA)] =
+      generator_registry[static_cast<int>(DeviceType::CUDA)] =
         detail::getCUDAHooks().initCUDAGenerator(this);
       detail::getCUDAHooks().registerCUDATypes(this);
     });
     return thc_state.get();
   }
+  THHState* lazyInitHIP() {
+    std::call_once(thh_init,[&] {
+      thh_state = detail::getHIPHooks().initHIP();
+      generator_registry[static_cast<int>(DeviceType::HIP)] =
+        detail::getHIPHooks().initHIPGenerator(this);
+      detail::getHIPHooks().registerHIPTypes(this);
+    });
+    return thh_state.get();
+  }
+  void lazyInitComplex() {
+    std::call_once(complex_init_, [&] {
+      detail::getComplexHooks().registerComplexTypes(this);
+    });
+  }
 
   THCState* getTHCState() {
     // AT_ASSERT(thc_state);
     return thc_state.get();
   }
-
-  CUDAStream createCUDAStream() const {
-    return detail::CUDAStream_createAndRetainWithOptions(
-      CUDAStream::DEFAULT_FLAGS
-    , CUDAStream::DEFAULT_PRIORITY
-    );
+  THHState* getTHHState() {
+    return thh_state.get();
   }
 
-  CUDAStream createCUDAStreamWithOptions(int32_t flags, int32_t priority) const {
-    return detail::CUDAStream_createAndRetainWithOptions(flags, priority);
-  }
-
-  CUDAStream getDefaultCUDAStream() const {
-    return detail::CUDAStream_getDefaultStream();
-  }
-
-  CUDAStream getDefaultCUDAStreamOnDevice(int64_t device) const {
-    return detail::CUDAStream_getDefaultStreamOnDevice(device);
-  }
-
-  CUDAStream getCurrentCUDAStream() const {
-    return detail::CUDAStream_getAndRetainCurrentStream();
-  }
-
-  CUDAStream getCurrentCUDAStreamOnDevice(int64_t device) const {
-    return detail::CUDAStream_getAndRetainCurrentStreamOnDevice(device);
-  }
-
-  void setCurrentCUDAStream(CUDAStream stream) const {
-    return detail::CUDAStream_setStream(stream.internals());
-  }
-
-  void setCurrentCUDAStreamOnDevice(int64_t device, CUDAStream stream) const {
-    return detail::CUDAStream_setStreamOnDevice(device, stream.internals());
-  }
-
-#ifndef __HIP_PLATFORM_HCC__
-  cusparseHandle_t getCurrentCUDASparseHandle() const {
-    return detail::getCUDAHooks().getCurrentCUDASparseHandle(thc_state.get());
-  }
-#endif
-  cudaDeviceProp* getCurrentDeviceProperties() const {
-    return detail::getCUDAHooks().getCurrentDeviceProperties(thc_state.get());
-  }
-  cudaDeviceProp* getDeviceProperties(int device) const {
-    return detail::getCUDAHooks().getDeviceProperties(thc_state.get(), device);
-  }
-  int getNumGPUs() const {
-    return detail::getCUDAHooks().getNumGPUs();
-  }
   size_t freshTypeID() {
     return next_id++;
   }
@@ -144,28 +125,36 @@ public:
   bool deterministicCuDNN() const;
   void setDeterministicCuDNN(bool);
   std::unique_ptr<Generator>
-    generator_registry[static_cast<int>(Backend::NumOptions)];
+    generator_registry[static_cast<int>(DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES)];
 private:
-  // NB: type_registry has nullptr for all CUDA backends until
-  // CUDA initialization has occurred
-  std::unique_ptr<Type> type_registry
-    [static_cast<int>(Backend::NumOptions)]
-    [static_cast<int>(ScalarType::NumOptions)];
-  void initCUDAIfNeeded(Backend p) {
-    if(p == Backend::CUDA)
+  void initCUDAIfNeeded(DeviceType p) {
+    if (p == DeviceType::CUDA) {
       lazyInitCUDA();
+    }
+  }
+  void initHIPIfNeeded(DeviceType p) {
+    if (p == DeviceType::HIP) {
+      lazyInitHIP();
+    }
+  }
+  void initComplexIfNeeded(ScalarType s) {
+    if (isComplexType(s)) {
+      lazyInitComplex();
+    }
   }
   std::once_flag thc_init;
+  std::once_flag thh_init;
+  std::once_flag complex_init_;
   bool enabled_cudnn = true;
   bool deterministic_cudnn = false;
   bool benchmark_cudnn = false;
   std::atomic<size_t> next_id;
   std::unique_ptr<THCState, void(*)(THCState*)> thc_state;
+  std::unique_ptr<THHState, void(*)(THHState*)> thh_state;
   friend struct Type;
-  friend void register_cuda_types(Context * context);
 };
 
-AT_API Context & globalContext();
+CAFFE2_API Context& globalContext();
 
 static inline void init() {
   globalContext();
@@ -177,32 +166,62 @@ static inline void init() {
   }
 }
 
-static inline Type& getType(Backend p, ScalarType s) {
-  return globalContext().getType(p, s);
+static inline TypeExtendedInterface& getNonVariableType(Backend p, ScalarType s) {
+  return globalContext().getNonVariableType(p, s);
 }
 
-static inline Type& CPU(ScalarType s) {
-  return getType(Backend::CPU, s);
+static inline TypeExtendedInterface& getNonVariableType(DeviceType p, ScalarType s) {
+  return globalContext().getNonVariableType(deviceTypeToBackend(p), s);
 }
 
-static inline Type& CUDA(ScalarType s) {
-  return getType(Backend::CUDA, s);
+CAFFE2_API TypeExtendedInterface& getType(TensorOptions options);
+CAFFE2_API TypeExtendedInterface& getType(const TensorImpl*);
+CAFFE2_API TypeExtendedInterface& getType(const Tensor&);
+
+CAFFE2_API Allocator* getCPUAllocator();
+
+static inline TypeExtendedInterface& CPU(ScalarType s) {
+  return getNonVariableType(Backend::CPU, s);
 }
 
+static inline TypeExtendedInterface& CUDA(ScalarType s) {
+  return getNonVariableType(Backend::CUDA, s);
+}
+
+static inline TypeExtendedInterface& HIP(ScalarType s) {
+  return getNonVariableType(Backend::HIP, s);
+}
+
+CAFFE2_API LegacyTHDispatcher& getLegacyTHDispatcher(TensorOptions options);
+CAFFE2_API LegacyTHDispatcher& getLegacyTHDispatcher(const Tensor&);
+
 static inline bool hasCUDA() {
   return globalContext().hasCUDA();
 }
 
-static inline bool hasCuDNN() {
-  return globalContext().hasCuDNN();
+static inline bool hasHIP() {
+  return globalContext().hasHIP();
 }
 
 static inline bool hasMKL() {
   return globalContext().hasMKL();
 }
 
-static inline int64_t current_device() {
-  return globalContext().current_device();
+static inline bool hasLAPACK() {
+  return globalContext().hasLAPACK();
+}
+
+static inline bool hasMAGMA() {
+  return globalContext().hasMAGMA();
+}
+
+static inline void manual_seed(uint64_t seed) {
+  globalContext().defaultGenerator(DeviceType::CPU).manualSeed(seed);
+  // NB: Sometimes we build with CUDA, but we don't have any GPUs
+  // available. In that case, we must not seed CUDA; it will fail!
+  if (hasCUDA() && detail::getCUDAHooks().getNumGPUs() > 0) {
+    globalContext().defaultGenerator(DeviceType::CUDA).manualSeedAll(seed);
+  }
 }
 
 } // namespace at

aten/src/ATen/DLConvertor.cpp

@@ -1,4 +1,5 @@
 #include "ATen/DLConvertor.h"
+#include "ATen/Functions.h"
 
 #include <iostream>
 #include <sstream>
@@ -36,6 +37,12 @@ static DLDataType getDLDataType(const Type& type) {
     case ScalarType::Half:
       dtype.code = DLDataTypeCode::kDLFloat;
       break;
+    case ScalarType::ComplexHalf:
+      throw std::logic_error("ComplexHalf is not supported by dlpack");
+    case ScalarType::ComplexFloat:
+      throw std::logic_error("ComplexFloat is not supported by dlpack");
+    case ScalarType::ComplexDouble:
+      throw std::logic_error("ComplexDouble is not supported by dlpack");
     case ScalarType::Undefined:
       throw std::logic_error("Undefined is not a valid ScalarType");
     case ScalarType::NumOptions:
@@ -57,19 +64,20 @@ static DLContext getDLContext(const Type& type, const int64_t& device_id) {
 }
 
 
-static Backend getATenBackend(const DLContext& ctx) {
-  Backend backend;
+static DeviceType getATenDeviceType(const DLContext& ctx) {
   switch (ctx.device_type) {
     case DLDeviceType::kDLCPU:
-      backend = Backend::CPU;
-      break;
+      return DeviceType::CPU;
     case DLDeviceType::kDLGPU:
-      backend = Backend::CUDA;
-      break;
+      return DeviceType::CUDA;
+    case DLDeviceType::kDLOpenCL:
+      return DeviceType::OPENCL;
+    case DLDeviceType::kDLROCM:
+      return DeviceType::HIP;
     default:
       throw std::logic_error("Unsupported device_type: " + std::to_string(ctx.device_type));
   }
-  return backend;
+  return DeviceType::CPU; // impossible
 }
 
 
@@ -144,7 +152,7 @@ DLManagedTensor* toDLPack(const Tensor& src) {
   atDLMTensor->tensor.deleter = &deleter;
   atDLMTensor->tensor.dl_tensor.data = src.data_ptr();
   int64_t device_id = 0;
-  if (src.type().is_cuda()) {
+  if (src.is_cuda()) {
     device_id = src.get_device();
   }
   atDLMTensor->tensor.dl_tensor.ctx = getDLContext(src.type(), device_id);
@@ -158,15 +166,15 @@ DLManagedTensor* toDLPack(const Tensor& src) {
 
 
 Tensor fromDLPack(const DLManagedTensor* src) {
-  Backend backend = getATenBackend(src->dl_tensor.ctx);
+  DeviceType device_type = getATenDeviceType(src->dl_tensor.ctx);
   ScalarType stype = toScalarType(src->dl_tensor.dtype);
   auto deleter = [src](void * self) {
     src->deleter(const_cast<DLManagedTensor*>(src));
   };
-  return getType(backend, stype).tensorFromBlob(
-      src->dl_tensor.data,
+  return at::from_blob(src->dl_tensor.data,
       IntList(src->dl_tensor.shape, src->dl_tensor.ndim),
       IntList(src->dl_tensor.strides, src->dl_tensor.ndim),
-      deleter);
+      deleter,
+      at::device(device_type).dtype(stype));
 }
 } //namespace at

aten/src/ATen/DLConvertor.h

@@ -10,8 +10,8 @@
 
 namespace at {
 
-AT_API ScalarType toScalarType(const DLDataType& dtype);
-AT_API DLManagedTensor * toDLPack(const Tensor& src);
-AT_API Tensor fromDLPack(const DLManagedTensor* src);
+CAFFE2_API ScalarType toScalarType(const DLDataType& dtype);
+CAFFE2_API DLManagedTensor* toDLPack(const Tensor& src);
+CAFFE2_API Tensor fromDLPack(const DLManagedTensor* src);
 
 } //namespace at

aten/src/ATen/Device.h

@@ -1,128 +1,2 @@
 #pragma once
-
-#include <ATen/Error.h>
-#include <ATen/ScalarType.h>
-
-#include <cstddef>
-#include <iosfwd>
-#include <string>
-#include <functional>
-
-namespace at {
-/// Represents a a compute device on which a tensor is located. A device is
-/// uniquely identified by a type, which specifies the type of machine it is
-/// (e.g. CPU or CUDA GPU), and a device index or ordinal, which identifies the
-/// specific compute device when there is more than one of a certain type. The
-/// device index is optional, and in its defaulted state represents (abstractly)
-/// "the current device". Further, there are two constraints on the value of the
-/// device index, if one is explicitly stored:
-/// 1. A negative index represents the current device, a non-negative index
-/// represents a specific, concrete device,
-/// 2. When the device type is CPU, the device index must be zero.
-struct Device {
-  /// The possible values of the device *type*.
-  enum class Type { CPU, CUDA };
-
-  /// Converts a `Backend` to a `Device::Type` if possible.
-  static Type backend_to_type(Backend backend) {
-    switch (backend) {
-      case kCPU:
-      case kSparseCPU:
-        return Type::CPU;
-      case kCUDA:
-      case kSparseCUDA:
-        return Type::CUDA;
-      default:
-        AT_ERROR(
-            "Invalid backend ", toString(backend), " for Device construction");
-    }
-  }
-
-  /// Constructs a new `Device` from a `Type` and an optional device index.
-  /* implicit */ Device(Type type, int32_t index = -1)
-      : type_(type), index_(index) {
-    AT_CHECK(
-        index == -1 || index >= 0,
-        "Device index must be -1 or non-negative, got ",
-        index);
-    AT_CHECK(
-        !is_cpu() || index <= 0,
-        "CPU device index must be -1 or zero, got ",
-        index);
-  }
-
-  /// Constructs a `Device` from a string description, for convenience.
-  /// The string supplied must follow the following schema:
-  /// `(cpu|cuda):[<device-index>]`
-  /// where `cpu:` or `cuda:` specifies the device type, and
-  /// `<device-index>` optionally specifies a device index.
-  /* implicit */ Device(const std::string& device_string);
-
-  /// Constructs a new `Device` from a `Backend` (which is converted to a
-  /// `Type`, if possible) and an optional device index.
-  /* implicit */ Device(Backend backend, int32_t index = -1)
-      : Device(backend_to_type(backend), index) {}
-
-  /// Returns true if the type and index of this `Device` matches that of
-  /// `other`.
-  bool operator==(const Device& other) const noexcept {
-    return this->type_ == other.type_ && this->index_ == other.index_;
-  }
-
-  /// Returns true if the type or index of this `Device` differs from that of
-  /// `other`.
-  bool operator!=(const Device& other) const noexcept {
-    return !(*this == other);
-  }
-
-  /// Sets the device index.
-  void set_index(int32_t index) {
-    index_ = index;
-  }
-
-  /// Returns the type of device this is.
-  Type type() const noexcept {
-    return type_;
-  }
-
-  /// Returns the optional index.
-  const int32_t& index() const noexcept {
-    return index_;
-  }
-
-  /// Returns true if the device has a non-default index.
-  bool has_index() const noexcept {
-    return index_ != -1;
-  }
-
-  /// Return true if the device is of CUDA type.
-  bool is_cuda() const noexcept {
-    return type_ == Type::CUDA;
-  }
-
-  /// Return true if the device is of CPU type.
-  bool is_cpu() const noexcept {
-    return type_ == Type::CPU;
-  }
-
- private:
-  Type type_;
-  int32_t index_ = -1;
-};
-} // namespace at
-
-std::ostream& operator<<(std::ostream& stream, at::Device::Type type);
-std::ostream& operator<<(std::ostream& stream, const at::Device& device);
-
-namespace std {
-  template<> struct hash<at::Device>
-  {
-    size_t operator()(const at::Device& device) const noexcept {
-      size_t hash_val = static_cast<size_t>(device.index() + 1);
-      if (device.is_cuda()) {
-        hash_val += 2;
-      }
-      return hash_val;
-    }
-  };
-} // namespace std
+#include <c10/Device.h>

aten/src/ATen/DeviceGuard.h

@@ -1,102 +1,36 @@
 #pragma once
 
-#include <ATen/Device.h>
-#include <ATen/Error.h>
-#include <ATen/ScalarType.h>
-#include <ATen/Tensor.h>
-#include <ATen/detail/CUDAHooksInterface.h>
-
-#include <cstddef>
+#include <c10/DeviceGuard.h>
+#include <ATen/core/Tensor.h>
+#include <c10/core/ScalarType.h> // TensorList whyyyyy
 
 namespace at {
-/// RAII guard that sets a certain default GPU index in its constructor, and
-/// changes it back to the device that was originally active upon destruction.
-///
-/// The index is always reset to the one that was active at the time of
-/// construction of the guard. Even if you `set_index` after construction, the
-/// destructor will still reset the index to the one that was active at
-/// construction time.
-struct DeviceGuard {
-  /// Default constructor, does nothing.
-  DeviceGuard() = default;
 
-  /// Uses the given device's `index()` if it is a CUDA device, else does
-  /// nothing.
-  explicit DeviceGuard(Device device) {
-    if (device.is_cuda()) {
-      set_index(device.index());
-    }
+// Are you here because you're wondering why DeviceGuard(tensor) no
+// longer works?  For code organization reasons, we have temporarily(?)
+// removed this constructor from DeviceGuard.  The new way to
+// spell it is:
+//
+//    OptionalDeviceGuard guard(device_of(tensor));
+
+/// Return the Device of a Tensor, if the Tensor is defined.
+inline optional<Device> device_of(Tensor t) {
+  if (t.defined()) {
+    return make_optional(t.device());
+  } else {
+    return nullopt;
   }
+}
 
-  /// Calls `set_device` with the given index.
-  explicit DeviceGuard(int32_t index) {
-    set_index(index);
+/// Return the Device of a TensorList, if the list is non-empty and
+/// the first Tensor is defined.  (This function implicitly assumes
+/// that all tensors in the list have the same device.)
+inline optional<Device> device_of(TensorList t) {
+  if (!t.empty()) {
+    return device_of(t.front());
+  } else {
+    return nullopt;
   }
+}
 
-  /// Sets the device to the index on which the given tensor is located.
-  explicit DeviceGuard(const Tensor& tensor) {
-    set_index_from(tensor);
-  }
-
-  /// Sets the device to the index on which the first tensor in the list is
-  /// located. If the list is empty, does nothing.
-  explicit DeviceGuard(const TensorList& tensors) {
-    if (!tensors.empty()) {
-      set_index_from(tensors.front());
-    }
-  }
-
-  /// Resets the device to the index that was active at construction of the
-  /// guard.
-  ~DeviceGuard() {
-    // It should only not have a value if an index was never actually set.
-    if (original_index_ != -1) {
-      // Unchecked because we don't want to throw in the destructor.
-      detail::DynamicCUDAInterface::unchecked_set_device(original_index_);
-    }
-  }
-
-  /// Sets the device to the given one.
-  void set_index(int32_t index) {
-    if (index == -1) {
-      return;
-    }
-    AT_ASSERT(index >= 0);
-    if (original_index_ == -1) {
-      int32_t previous_index = -123;
-      detail::DynamicCUDAInterface::get_device(&previous_index);
-      original_index_ = previous_index;
-      if (index != original_index_) {
-        detail::DynamicCUDAInterface::set_device(index);
-      }
-    } else {
-      detail::DynamicCUDAInterface::set_device(index);
-    }
-    last_index_ = index;
-  }
-
-  /// Calls `set_index` with the `Tensor`'s current device, if it is a CUDA
-  /// tensor. Does nothing if the `tensor` is not defined.
-  void set_index_from(const Tensor& tensor) {
-    if (tensor.defined() && tensor.is_cuda()) {
-      set_index(tensor.get_device());
-    }
-  }
-
-  /// Returns the device that was set upon construction of the guard.
-  int32_t original_index() const noexcept {
-    return original_index_;
-  }
-
-  // /// Returns the last device that was set via `set_device`, if any.
-  int32_t last_index() const noexcept {
-    return last_index_;
-  }
-
- private:
-  /// The original device that was active at construction of this object.
-  int32_t original_index_ = -1;
-  /// The last index that was set via `set_device`.
-  int32_t last_index_ = -1;
-};
 } // namespace at

aten/src/ATen/DimVector.h

@@ -1,11 +1,2 @@
 #pragma once
-
-#include "SmallVector.h"
-#include <stdint.h>
-
-namespace at {
-
-/// A container for sizes or strides
-using DimVector = SmallVector<int64_t, 5>;
-
-}
+#include <ATen/core/DimVector.h>

aten/src/ATen/Dispatch.h

@@ -1,8 +1,8 @@
 #pragma once
 
-#include <ATen/Error.h>
-#include <ATen/Half.h>
 #include <ATen/Type.h>
+#include <ATen/core/Half.h>
+#include <c10/util/Exception.h>
 
 #define AT_PRIVATE_CASE_TYPE(enum_type, type, ...) \
   case enum_type: {                                \
@@ -10,72 +10,144 @@
     return __VA_ARGS__();                          \
   }
 
-#define AT_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...)                           \
-  [&] {                                                                       \
-    const at::Type& the_type = TYPE;                                          \
-    switch (the_type.scalarType()) {                                          \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Double, double, __VA_ARGS__)       \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Float, float, __VA_ARGS__)         \
-      default:                                                                \
+#define AT_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...)                          \
+  [&] {                                                                      \
+    const at::Type& the_type = TYPE;                                         \
+    switch (the_type.scalarType()) {                                         \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Double, double, __VA_ARGS__)      \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Float, float, __VA_ARGS__)        \
+      default:                                                               \
         AT_ERROR(#NAME, " not implemented for '", the_type.toString(), "'"); \
-    }                                                                         \
+    }                                                                        \
   }()
 
-#define AT_DISPATCH_FLOATING_TYPES_AND_HALF(TYPE, NAME, ...)                  \
-  [&] {                                                                       \
-    const at::Type& the_type = TYPE;                                          \
-    switch (the_type.scalarType()) {                                          \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Double, double, __VA_ARGS__)       \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Float, float, __VA_ARGS__)         \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Half, Half, __VA_ARGS__)           \
-      default:                                                                \
+#define AT_DISPATCH_FLOATING_TYPES_AND_HALF(TYPE, NAME, ...)                 \
+  [&] {                                                                      \
+    const at::Type& the_type = TYPE;                                         \
+    switch (the_type.scalarType()) {                                         \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Double, double, __VA_ARGS__)      \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Float, float, __VA_ARGS__)        \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Half, at::Half, __VA_ARGS__)      \
+      default:                                                               \
         AT_ERROR(#NAME, " not implemented for '", the_type.toString(), "'"); \
-    }                                                                         \
+    }                                                                        \
   }()
 
-#define AT_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...)                           \
-  [&] {                                                                       \
-    const at::Type& the_type = TYPE;                                          \
-    switch (the_type.scalarType()) {                                          \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Byte, uint8_t, __VA_ARGS__)        \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Char, int8_t, __VA_ARGS__)         \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Int, int32_t, __VA_ARGS__)         \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Long, int64_t, __VA_ARGS__)        \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Short, int16_t, __VA_ARGS__)       \
-      default:                                                                \
-        AT_ERROR("%s not implemented for '%s'", (NAME), the_type.toString()); \
-    }                                                                         \
+#define AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(TYPE, NAME, ...)              \
+  [&] {                                                                      \
+    const at::Type& the_type = TYPE;                                         \
+    switch (the_type.scalarType()) {                                         \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Double, double, __VA_ARGS__)      \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Float, float, __VA_ARGS__)        \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Half, at::Half, __VA_ARGS__)      \
+      AT_PRIVATE_CASE_TYPE(                                                  \
+          at::ScalarType::ComplexDouble, std::complex<double>, __VA_ARGS__)  \
+      AT_PRIVATE_CASE_TYPE(                                                  \
+          at::ScalarType::ComplexFloat, std::complex<float>, __VA_ARGS__)    \
+      AT_PRIVATE_CASE_TYPE(                                                  \
+          at::ScalarType::ComplexHalf, std::complex<at::Half>, __VA_ARGS__)  \
+      default:                                                               \
+        AT_ERROR(#NAME, " not implemented for '", the_type.toString(), "'"); \
+    }                                                                        \
   }()
 
-#define AT_DISPATCH_ALL_TYPES(TYPE, NAME, ...)                                \
-  [&] {                                                                       \
-    const at::Type& the_type = TYPE;                                          \
-    switch (the_type.scalarType()) {                                          \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Byte, uint8_t, __VA_ARGS__)        \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Char, int8_t, __VA_ARGS__)         \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Double, double, __VA_ARGS__)       \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Float, float, __VA_ARGS__)         \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Int, int32_t, __VA_ARGS__)         \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Long, int64_t, __VA_ARGS__)        \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Short, int16_t, __VA_ARGS__)       \
-      default:                                                                \
+#define AT_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...)                          \
+  [&] {                                                                      \
+    const at::Type& the_type = TYPE;                                         \
+    switch (the_type.scalarType()) {                                         \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Byte, uint8_t, __VA_ARGS__)       \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Char, int8_t, __VA_ARGS__)        \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Int, int32_t, __VA_ARGS__)        \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Long, int64_t, __VA_ARGS__)       \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Short, int16_t, __VA_ARGS__)      \
+      default:                                                               \
         AT_ERROR(#NAME, " not implemented for '", the_type.toString(), "'"); \
-    }                                                                         \
+    }                                                                        \
   }()
 
-#define AT_DISPATCH_ALL_TYPES_AND_HALF(TYPE, NAME, ...)                       \
-  [&] {                                                                       \
-    const at::Type& the_type = TYPE;                                          \
-    switch (the_type.scalarType()) {                                          \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Byte, uint8_t, __VA_ARGS__)        \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Char, int8_t, __VA_ARGS__)         \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Double, double, __VA_ARGS__)       \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Float, float, __VA_ARGS__)         \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Int, int32_t, __VA_ARGS__)         \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Long, int64_t, __VA_ARGS__)        \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Short, int16_t, __VA_ARGS__)       \
-      AT_PRIVATE_CASE_TYPE(at::ScalarType::Half, Half, __VA_ARGS__)           \
-      default:                                                                \
+#define AT_DISPATCH_ALL_TYPES(TYPE, NAME, ...)                               \
+  [&] {                                                                      \
+    const at::Type& the_type = TYPE;                                         \
+    switch (the_type.scalarType()) {                                         \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Byte, uint8_t, __VA_ARGS__)       \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Char, int8_t, __VA_ARGS__)        \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Double, double, __VA_ARGS__)      \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Float, float, __VA_ARGS__)        \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Int, int32_t, __VA_ARGS__)        \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Long, int64_t, __VA_ARGS__)       \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Short, int16_t, __VA_ARGS__)      \
+      default:                                                               \
         AT_ERROR(#NAME, " not implemented for '", the_type.toString(), "'"); \
-    }                                                                         \
+    }                                                                        \
+  }()
+
+#define AT_DISPATCH_ALL_TYPES_AND_HALF(TYPE, NAME, ...)                      \
+  [&] {                                                                      \
+    const at::Type& the_type = TYPE;                                         \
+    switch (the_type.scalarType()) {                                         \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Byte, uint8_t, __VA_ARGS__)       \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Char, int8_t, __VA_ARGS__)        \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Double, double, __VA_ARGS__)      \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Float, float, __VA_ARGS__)        \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Int, int32_t, __VA_ARGS__)        \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Long, int64_t, __VA_ARGS__)       \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Short, int16_t, __VA_ARGS__)      \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Half, at::Half, __VA_ARGS__)      \
+      default:                                                               \
+        AT_ERROR(#NAME, " not implemented for '", the_type.toString(), "'"); \
+    }                                                                        \
+  }()
+
+#define AT_DISPATCH_COMPLEX_TYPES(TYPE, NAME, ...)                           \
+  [&] {                                                                      \
+    const at::Type& the_type = TYPE;                                         \
+    switch (the_type.scalarType()) {                                         \
+      AT_PRIVATE_CASE_TYPE(                                                  \
+          at::ScalarType::ComplexFloat, std::complex<float>, __VA_ARGS__)    \
+      AT_PRIVATE_CASE_TYPE(                                                  \
+          at::ScalarType::ComplexDouble, std::complex<double>, __VA_ARGS__)  \
+      default:                                                               \
+        AT_ERROR(#NAME, " not implemented for '", the_type.toString(), "'"); \
+    }                                                                        \
+  }()
+
+#define AT_DISPATCH_ALL_TYPES_AND_COMPLEX(TYPE, NAME, ...)                   \
+  [&] {                                                                      \
+    const at::Type& the_type = TYPE;                                         \
+    switch (the_type.scalarType()) {                                         \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Byte, uint8_t, __VA_ARGS__)       \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Char, int8_t, __VA_ARGS__)        \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Double, double, __VA_ARGS__)      \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Float, float, __VA_ARGS__)        \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Int, int32_t, __VA_ARGS__)        \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Long, int64_t, __VA_ARGS__)       \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Short, int16_t, __VA_ARGS__)      \
+      AT_PRIVATE_CASE_TYPE(                                                  \
+          at::ScalarType::ComplexFloat, std::complex<float>, __VA_ARGS__)    \
+      AT_PRIVATE_CASE_TYPE(                                                  \
+          at::ScalarType::ComplexDouble, std::complex<double>, __VA_ARGS__)  \
+      default:                                                               \
+        AT_ERROR(#NAME, " not implemented for '", the_type.toString(), "'"); \
+    }                                                                        \
+  }()
+
+#define AT_DISPATCH_ALL_TYPES_AND_HALF_AND_COMPLEX(TYPE, NAME, ...)          \
+  [&] {                                                                      \
+    const at::Type& the_type = TYPE;                                         \
+    switch (the_type.scalarType()) {                                         \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Byte, uint8_t, __VA_ARGS__)       \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Char, int8_t, __VA_ARGS__)        \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Double, double, __VA_ARGS__)      \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Float, float, __VA_ARGS__)        \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Int, int32_t, __VA_ARGS__)        \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Long, int64_t, __VA_ARGS__)       \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Short, int16_t, __VA_ARGS__)      \
+      AT_PRIVATE_CASE_TYPE(at::ScalarType::Half, at::Half, __VA_ARGS__)      \
+      AT_PRIVATE_CASE_TYPE(                                                  \
+          at::ScalarType::ComplexFloat, std::complex<float>, __VA_ARGS__)    \
+      AT_PRIVATE_CASE_TYPE(                                                  \
+          at::ScalarType::ComplexDouble, std::complex<double>, __VA_ARGS__)  \
+      default:                                                               \
+        AT_ERROR(#NAME, " not implemented for '", the_type.toString(), "'"); \
+    }                                                                        \
   }()

aten/src/ATen/Error.cpp

@@ -1,32 +0,0 @@
-#include <ATen/Error.h>
-#include <ATen/Backtrace.h>
-
-#include <iostream>
-#include <string>
-
-namespace at {
-std::ostream& operator<<(std::ostream& out, const SourceLocation& loc) {
-  out << loc.function << " at " << loc.file << ":" << loc.line;
-  return out;
-}
-
-Error::Error(SourceLocation source_location, std::string err)
-  : what_without_backtrace_(err)
-  , what_(str(err, " (", source_location, ")\n", get_backtrace(/*frames_to_skip=*/2)))
-  {}
-
-void Warning::warn(SourceLocation source_location, std::string msg) {
-  warning_handler_(source_location, msg.c_str());
-}
-
-void Warning::set_warning_handler(handler_t handler) {
-  warning_handler_ = handler;
-}
-
-void Warning::print_warning(const SourceLocation& source_location, const char* msg) {
-  std::cerr << "Warning: " << msg << " (" << source_location << ")\n";
-}
-
-Warning::handler_t Warning::warning_handler_ = &Warning::print_warning;
-
-} // namespace at

aten/src/ATen/Error.h

@@ -1,131 +0,0 @@
-#pragma once
-
-#include <ATen/ATenGeneral.h> // for AT_API
-#include <ATen/optional.h>
-
-#include <cstddef>
-#include <exception>
-#include <ostream>
-#include <sstream>
-#include <string>
-
-#if defined(_MSC_VER) && _MSC_VER <= 1900
-#define __func__ __FUNCTION__
-#endif
-
-namespace at {
-
-namespace detail {
-
-inline std::ostream& _str(std::ostream& ss) { return ss; }
-
-template <typename T>
-inline std::ostream& _str(std::ostream& ss, const T& t) {
-  ss << t;
-  return ss;
-}
-
-template <typename T, typename... Args>
-inline std::ostream&
-_str(std::ostream& ss, const T& t, const Args&... args) {
-  return _str(_str(ss, t), args...);
-}
-
-} // namespace detail
-
-// Convert a list of string-like arguments into a single string.
-template <typename... Args>
-inline std::string str(const Args&... args) {
-  std::ostringstream ss;
-  detail::_str(ss, args...);
-  return ss.str();
-}
-
-// Specializations for already-a-string types.
-template <>
-inline std::string str(const std::string& str) {
-  return str;
-}
-inline std::string str(const char* c_str) {
-  return c_str;
-}
-
-/// Represents a location in source code (for debugging).
-struct SourceLocation {
-  const char* function;
-  const char* file;
-  uint32_t line;
-};
-
-std::ostream& operator<<(std::ostream& out, const SourceLocation& loc);
-
-/// The primary ATen error class.
-/// Provides a complete error message with source location information via
-/// `what()`, and a more concise message via `what_without_backtrace()`. Should
-/// primarily be used with the `AT_ERROR` macro.
-///
-/// NB: at::Error is handled specially by the default torch to suppress the
-/// backtrace, see torch/csrc/Exceptions.h
-class AT_API Error : public std::exception {
-  std::string what_without_backtrace_;
-  std::string what_;
-
-public:
-  Error(SourceLocation source_location, std::string err);
-
-  /// Returns the complete error message, including the source location.
-  const char* what() const noexcept override {
-    return what_.c_str();
-  }
-
-  /// Returns only the error message string, without source location.
-  const char* what_without_backtrace() const noexcept {
-    return what_without_backtrace_.c_str();
-  }
-};
-
-class AT_API Warning {
-  using handler_t = void(*)(const SourceLocation& source_location, const char* msg);
-
-public:
-  /// Issue a warning with a given message. Dispatched to the current
-  /// warning handler.
-  static void warn(SourceLocation source_location, std::string msg);
-
-  /// Sets the global warning handler. This is not thread-safe, so it should
-  /// generally be called once during initialization.
-  static void set_warning_handler(handler_t handler);
-
-  /// The default warning handler. Prints the message to stderr.
-  static void print_warning(const SourceLocation& source_location, const char* msg);
-
-private:
-  static handler_t warning_handler_;
-};
-
-
-} // namespace at
-
-// TODO: variants that print the expression tested and thus don't require strings
-// TODO: CAFFE_ENFORCE_WITH_CALLER style macro
-
-#define AT_ERROR(...) \
-  throw at::Error({__func__, __FILE__, __LINE__}, at::str(__VA_ARGS__))
-
-#define AT_WARN(...) \
-  at::Warning::warn({__func__, __FILE__, __LINE__}, at::str(__VA_ARGS__))
-
-#define AT_ASSERT(cond) \
-  if (!(cond)) {             \
-    AT_ERROR(#cond " ASSERT FAILED at ", __FILE__, ":", __LINE__, ", please report a bug to PyTorch.");   \
-  }
-
-#define AT_ASSERTM(cond, ...) \
-  if (!(cond)) {             \
-    AT_ERROR(at::str(#cond, " ASSERT FAILED at ", __FILE__, ":", __LINE__, ", please report a bug to PyTorch. ", __VA_ARGS__));   \
-  }
-
-#define AT_CHECK(cond, ...) \
-  if (!(cond)) {             \
-    AT_ERROR(at::str(__VA_ARGS__));   \
-  }

aten/src/ATen/ExpandUtils.cpp

@@ -29,11 +29,13 @@ std::vector<int64_t> infer_size(IntList a, IntList b) {
 }
 
 std::tuple<std::vector<int64_t>, std::vector<int64_t>> inferExpandGeometry(
-    const Tensor& tensor,
+    IntList tensor_sizes,
+    IntList tensor_strides,
     IntList sizes) {
   int64_t ndim = sizes.size();
+  int64_t tensor_dim = tensor_sizes.size();
 
-  if (tensor.dim() == 0) {
+  if (tensor_dim == 0) {
     std::vector<int64_t> expandedStrides(ndim, 0);
     return std::tuple<std::vector<int64_t>, std::vector<int64_t>>(
         sizes.vec(), expandedStrides);
@@ -44,9 +46,9 @@ std::tuple<std::vector<int64_t>, std::vector<int64_t>> inferExpandGeometry(
   // create a new geometry for the tensors
   for (int64_t i = ndim - 1; i >= 0; --i) {
     int64_t offset = ndim - 1 - i;
-    int64_t dim = tensor.dim() - 1 - offset;
-    int64_t size = (dim >= 0) ? tensor.sizes()[dim] : 1;
-    int64_t stride = (dim >= 0) ? tensor.strides()[dim]
+    int64_t dim = tensor_dim - 1 - offset;
+    int64_t size = (dim >= 0) ? tensor_sizes[dim] : 1;
+    int64_t stride = (dim >= 0) ? tensor_strides[dim]
                                 : expandedSizes[i + 1] * expandedStrides[i + 1];
     int64_t targetSize = sizes[i];
     if (targetSize == -1) {
@@ -66,7 +68,11 @@ std::tuple<std::vector<int64_t>, std::vector<int64_t>> inferExpandGeometry(
           ") must match the existing size (",
           size,
           ") at non-singleton dimension ",
-          i);
+          i,
+          ".  Target sizes: ",
+          sizes,
+          ".  Tensor sizes: ",
+          tensor_sizes);
       size = targetSize;
       stride = 0;
     }

aten/src/ATen/ExpandUtils.h

@@ -1,7 +1,7 @@
 #pragma once
 
 #include "ATen/Tensor.h"
-#include "ATen/Error.h"
+#include "c10/util/Exception.h"
 
 #include <functional>
 #include <sstream>
@@ -9,8 +9,12 @@
 
 namespace at {
 
-AT_API std::vector<int64_t> infer_size(IntList a, IntList b);
-std::tuple<std::vector<int64_t>, std::vector<int64_t> > inferExpandGeometry(const Tensor &tensor, IntList sizes);
+CAFFE2_API std::vector<int64_t> infer_size(IntList a, IntList b);
+CAFFE2_API std::tuple<std::vector<int64_t>, std::vector<int64_t>>
+inferExpandGeometry(
+    IntList tensor_sizes,
+    IntList tensor_strides,
+    IntList sizes);
 
 // avoid copy-construction of Tensor by using a reference_wrapper.
 inline void check_defined(std::initializer_list<std::reference_wrapper<const Tensor>> tensors, const char *api_name) {
@@ -110,7 +114,7 @@ inline std::vector<Tensor> expand_outplace(TensorList to_expand) {
     if (!to_expand[i].defined()) {
       continue;
     } else if (first) {
-      sizes = to_expand[i].sizes();
+      sizes = to_expand[i].sizes().vec();
       first = false;
     } else {
       sizes = infer_size(sizes, to_expand[i].sizes());
@@ -130,4 +134,44 @@ inline std::vector<Tensor> expand_outplace(TensorList to_expand) {
   return result;
 }
 
+// Sums `tensor` repeatedly to produce a tensor of shape `shape`.
+// Precondition: is_expandable_to(shape, tensor.sizes()) must be true
+static inline Tensor sum_to(Tensor tensor, const IntList shape) {
+  if (shape.size() == 0) {
+    return tensor.sum();
+  }
+  c10::SmallVector<int64_t, 8> reduce_dims;
+  const at::IntList sizes = tensor.sizes();
+  const int64_t leading_dims = sizes.size() - shape.size();
+  for (int64_t i = 0; i < leading_dims; ++i) {
+    reduce_dims.push_back(i);
+  }
+  for (int64_t i = leading_dims; i < static_cast<int64_t>(sizes.size()); ++i) {
+    if (shape[i - leading_dims] == 1 && sizes[i] > 1) {
+      reduce_dims.push_back(i);
+    }
+  }
+  if (!reduce_dims.empty()) {
+    tensor = tensor.sum(reduce_dims, /*keepdim=*/true);
+  }
+  return leading_dims > 0 ? tensor.view(shape) : tensor;
+}
+
+// True if `shape` can be broadcasted to `desired`
+static inline bool is_expandable_to(IntList shape, IntList desired) {
+  int ndim = shape.size();
+  int target_dim = desired.size();
+  if (ndim > target_dim) {
+    return false;
+  }
+  for (int i = 0; i < ndim; i++) {
+    int64_t size = shape[ndim - i - 1];
+    int64_t target = desired[target_dim - i - 1];
+    if (size != target && size != 1) {
+      return false;
+    }
+  }
+  return true;
+}
+
 }

aten/src/ATen/Formatting.h

@@ -1,26 +1 @@
-#pragma once
-
-#include <iostream>
-#include "ATen/Type.h"
-#include "ATen/Scalar.h"
-
-namespace at {
-
-AT_API std::ostream& operator<<(std::ostream & out, IntList list);
-AT_API std::ostream& operator<<(std::ostream & out, Backend b);
-AT_API std::ostream& operator<<(std::ostream & out, ScalarType t);
-AT_API std::ostream& operator<<(std::ostream & out, const Type & t);
-AT_API std::ostream& print(std::ostream& stream, const Tensor & tensor, int64_t linesize);
-static inline std::ostream& operator<<(std::ostream & out, const Tensor & t) {
-  return print(out,t,80);
-}
-static inline void print(const Tensor & t, int64_t linesize=80) {
-  print(std::cout,t,linesize);
-}
-
-static inline std::ostream& operator<<(std::ostream & out, Scalar s) {
-  s = s.local();
-  return out << (s.isFloatingPoint() ? s.toDouble() : s.toLong());
-}
-
-}
+#include <ATen/core/Formatting.h>

aten/src/ATen/Generator.h

@@ -1,23 +1,2 @@
 #pragma once
-
-#include <stdint.h>
-
-namespace at {
-
-struct Generator {
-  Generator() {};
-  Generator(const Generator& other) = delete;
-  Generator(Generator&& other) = delete;
-  virtual ~Generator() {};
-
-  virtual Generator& copy(const Generator& other) = 0;
-  virtual Generator& free() = 0;
-
-  virtual uint64_t seed() = 0;
-  virtual uint64_t initialSeed() = 0;
-  virtual Generator& manualSeed(uint64_t seed) = 0;
-  virtual Generator& manualSeedAll(uint64_t seed) = 0;
-  virtual void * unsafeGetTH() = 0;
-};
-
-} // namespace at
+#include <ATen/core/Generator.h>

aten/src/ATen/Half-inl.h

@@ -1,168 +0,0 @@
-#pragma once
-
-#include "ATen/ATenGeneral.h"
-#include <cstring>
-#include <limits>
-
-#ifdef __CUDACC__
-#include <cuda_fp16.h>
-#endif
-
-namespace at {
-
-/// Constructors
-
-inline AT_HOSTDEVICE Half::Half(float value) {
-#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
-  x = __half_as_short(__float2half(value));
-#else
-  x = detail::float2halfbits(value);
-#endif
-}
-
-/// Implicit conversions
-
-inline AT_HOSTDEVICE Half::operator float() const {
-#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
-  return __half2float(*reinterpret_cast<const __half*>(&x));
-#else
-  return detail::halfbits2float(x);
-#endif
-}
-
-#ifdef __CUDACC__
-inline AT_HOSTDEVICE Half::Half(const __half& value) {
-  x = *reinterpret_cast<const unsigned short*>(&value);
-}
-inline AT_HOSTDEVICE Half::operator __half() const {
-  return *reinterpret_cast<const __half*>(&x);
-}
-#endif
-
-/// Arithmetic
-
-inline AT_HOSTDEVICE Half operator+(const Half& a, const Half& b) {
-  return (float)a + (float)b;
-}
-
-inline AT_HOSTDEVICE Half operator-(const Half& a, const Half& b) {
-  return (float)a - (float)b;
-}
-
-inline AT_HOSTDEVICE Half operator*(const Half& a, const Half& b) {
-  return (float)a * (float)b;
-}
-
-inline AT_HOSTDEVICE Half operator/(const Half& a, const Half& b) {
-  return (float)a / (float)b;
-}
-
-inline AT_HOSTDEVICE Half operator-(const Half& a) {
-  return -(float)a;
-}
-
-inline AT_HOSTDEVICE Half& operator+=(Half& a, const Half& b) {
-  a = a + b;
-  return a;
-}
-
-inline AT_HOSTDEVICE Half& operator-=(Half& a, const Half& b) {
-  a = a - b;
-  return a;
-}
-
-inline AT_HOSTDEVICE Half& operator*=(Half& a, const Half& b) {
-  a = a * b;
-  return a;
-}
-
-inline AT_HOSTDEVICE Half& operator/=(Half& a, const Half& b) {
-  a = a / b;
-  return a;
-}
-
-/// Arithmetic with floats
-
-inline AT_HOSTDEVICE float operator+(Half a, float b) { return (float)a + b; }
-inline AT_HOSTDEVICE float operator-(Half a, float b) { return (float)a - b; }
-inline AT_HOSTDEVICE float operator*(Half a, float b) { return (float)a * b; }
-inline AT_HOSTDEVICE float operator/(Half a, float b) { return (float)a / b; }
-
-inline AT_HOSTDEVICE float operator+(float a, Half b) { return a + (float)b; }
-inline AT_HOSTDEVICE float operator-(float a, Half b) { return a - (float)b; }
-inline AT_HOSTDEVICE float operator*(float a, Half b) { return a * (float)b; }
-inline AT_HOSTDEVICE float operator/(float a, Half b) { return a / (float)b; }
-
-inline AT_HOSTDEVICE float& operator+=(float& a, const Half& b) { return a += (float)b; }
-inline AT_HOSTDEVICE float& operator-=(float& a, const Half& b) { return a -= (float)b; }
-inline AT_HOSTDEVICE float& operator*=(float& a, const Half& b) { return a *= (float)b; }
-inline AT_HOSTDEVICE float& operator/=(float& a, const Half& b) { return a /= (float)b; }
-
-/// Arithmetic with doubles
-
-inline AT_HOSTDEVICE double operator+(Half a, double b) { return (double)a + b; }
-inline AT_HOSTDEVICE double operator-(Half a, double b) { return (double)a - b; }
-inline AT_HOSTDEVICE double operator*(Half a, double b) { return (double)a * b; }
-inline AT_HOSTDEVICE double operator/(Half a, double b) { return (double)a / b; }
-
-inline AT_HOSTDEVICE double operator+(double a, Half b) { return a + (double)b; }
-inline AT_HOSTDEVICE double operator-(double a, Half b) { return a - (double)b; }
-inline AT_HOSTDEVICE double operator*(double a, Half b) { return a * (double)b; }
-inline AT_HOSTDEVICE double operator/(double a, Half b) { return a / (double)b; }
-
-/// Arithmetic with ints
-
-inline AT_HOSTDEVICE Half operator+(Half a, int b) { return a + (Half)b; }
-inline AT_HOSTDEVICE Half operator-(Half a, int b) { return a - (Half)b; }
-inline AT_HOSTDEVICE Half operator*(Half a, int b) { return a * (Half)b; }
-inline AT_HOSTDEVICE Half operator/(Half a, int b) { return a / (Half)b; }
-
-inline AT_HOSTDEVICE Half operator+(int a, Half b) { return (Half)a + b; }
-inline AT_HOSTDEVICE Half operator-(int a, Half b) { return (Half)a - b; }
-inline AT_HOSTDEVICE Half operator*(int a, Half b) { return (Half)a * b; }
-inline AT_HOSTDEVICE Half operator/(int a, Half b) { return (Half)a / b; }
-
-/// NOTE: we do not define comparisons directly and instead rely on the implicit
-/// conversion from at::Half to float.
-
-} // namespace at
-
-namespace std {
-
-template<> class numeric_limits<at::Half> {
- public:
-  static constexpr bool is_specialized = true;
-  static constexpr bool is_signed = true;
-  static constexpr bool is_integer = false;
-  static constexpr bool is_exact = false;
-  static constexpr bool has_infinity = true;
-  static constexpr bool has_quiet_NaN = true;
-  static constexpr bool has_signaling_NaN = true;
-  static constexpr auto has_denorm = numeric_limits<float>::has_denorm;
-  static constexpr auto has_denorm_loss = numeric_limits<float>::has_denorm_loss;
-  static constexpr auto round_style = numeric_limits<float>::round_style;
-  static constexpr bool is_iec559 = true;
-  static constexpr bool is_bounded = true;
-  static constexpr bool is_modulo = false;
-  static constexpr int digits = 11;
-  static constexpr int digits10 = 3;
-  static constexpr int max_digits10 = 5;
-  static constexpr int radix = 2;
-  static constexpr int min_exponent = -13;
-  static constexpr int min_exponent10 = -4;
-  static constexpr int max_exponent = 16;
-  static constexpr int max_exponent10 = 4;
-  static constexpr auto traps = numeric_limits<float>::traps;
-  static constexpr auto tinyness_before = numeric_limits<float>::tinyness_before;
-  static constexpr at::Half min() { return at::Half(0x0400, at::Half::from_bits); }
-  static constexpr at::Half lowest() { return at::Half(0xFBFF, at::Half::from_bits); }
-  static constexpr at::Half max() { return at::Half(0x7BFF, at::Half::from_bits); }
-  static constexpr at::Half epsilon() { return at::Half(0x1400, at::Half::from_bits); }
-  static constexpr at::Half round_error() { return at::Half(0x3800, at::Half::from_bits); }
-  static constexpr at::Half infinity() { return at::Half(0x7C00, at::Half::from_bits); }
-  static constexpr at::Half quiet_NaN() { return at::Half(0x7E00, at::Half::from_bits); }
-  static constexpr at::Half signaling_NaN() { return at::Half(0x7D00, at::Half::from_bits); }
-  static constexpr at::Half denorm_min() { return at::Half(0x0001, at::Half::from_bits); }
-};
-
-} // namespace std

aten/src/ATen/Half.cpp

@@ -1,34 +0,0 @@
-#include "ATen/Half.h"
-
-#include "ATen/Tensor.h"
-#include "ATen/Context.h"
-
-#include <TH/TH.h>
-#include <iostream>
-
-namespace at {
-
-static_assert(std::is_standard_layout<Half>::value, "at::Half must be standard layout.");
-
-namespace detail {
-
-float halfbits2float(unsigned short bits) {
-  float value;
-  TH_halfbits2float(&bits, &value);
-  return value;
-}
-
-unsigned short float2halfbits(float value) {
-  unsigned short bits;
-  TH_float2halfbits(&value, &bits);
-  return bits;
-}
-
-} // namespace detail
-
-std::ostream& operator<<(std::ostream & out, const Half& value) {
-  out << (float)value;
-  return out;
-}
-
-} // namespace at

aten/src/ATen/Half.h

@@ -1,113 +1,2 @@
 #pragma once
-
-/// Defines the Half type (half-precision floating-point) including conversions
-/// to standard C types and basic arithmetic operations. Note that arithmetic
-/// operations are implemented by converting to floating point and
-/// performing the operation in float32, instead of using CUDA half intrinisics.
-/// Most uses of this type within ATen are memory bound, including the
-/// element-wise kernels, and the half intrinisics aren't efficient on all GPUs.
-/// If you are writing a compute bound kernel, you can use the CUDA half
-/// intrinsics directly on the Half type from device code.
-
-#include "ATen/ATenGeneral.h"
-
-#include <limits>
-#include <string>
-#include <cstdint>
-#include <stdexcept>
-#include <utility>
-#include <cmath>
-#include <iosfwd>
-
-#ifdef __CUDACC__
-#include <cuda_fp16.h>
-#endif
-
-#ifndef AT_HOSTDEVICE
-  #ifdef __CUDACC__
-    #define AT_HOSTDEVICE __host__ __device__
-  #else
-    #define AT_HOSTDEVICE
-  #endif
-#endif
-
-namespace at {
-
-namespace detail {
-
-float halfbits2float(unsigned short bits);
-unsigned short float2halfbits(float value);
-
-}
-
-struct alignas(2) Half {
-  unsigned short x;
-
-  struct from_bits_t {};
-  static constexpr from_bits_t from_bits = from_bits_t();
-
-  // HIP wants __host__ __device__ tag, CUDA does not
-#ifdef __HIP_PLATFORM_HCC__
-  AT_HOSTDEVICE Half() = default;
-#else
-  Half() = default;
-#endif
-
-  constexpr AT_HOSTDEVICE Half(unsigned short bits, from_bits_t) : x(bits) {};
-  inline AT_HOSTDEVICE Half(float value);
-  inline AT_HOSTDEVICE operator float() const;
-
-#ifdef __CUDACC__
-  inline AT_HOSTDEVICE Half(const __half& value);
-  inline AT_HOSTDEVICE operator __half() const;
-#endif
-};
-
-template<typename To, typename From> To convert(From f) {
-  return static_cast<To>(f);
-}
-
-// skip isnan and isinf check for integral types
-template<typename To, typename From>
-typename std::enable_if<std::is_integral<From>::value, bool>::type overflows(From f) {
-  using limit = std::numeric_limits<To>;
-  return f < limit::lowest() || f > limit::max();
-}
-
-template<typename To, typename From>
-typename std::enable_if<!std::is_integral<From>::value, bool>::type overflows(From f) {
-  using limit = std::numeric_limits<To>;
-  if (limit::has_infinity && std::isinf((double)f)) {
-    return false;
-  }
-  if (!limit::has_quiet_NaN && (f != f)) {
-    return true;
-  }
-  return f < limit::lowest() || f > limit::max();
-}
-
-template<typename To, typename From> To checked_convert(From f, const char* name) {
-  if (overflows<To, From>(f)) {
-    std::string msg = "value cannot be converted to type ";
-    msg += name;
-    msg += " without overflow: ";
-    msg += std::to_string(f);
-    throw std::domain_error(std::move(msg));
-  }
-  return convert<To, From>(f);
-}
-
-template<typename To, typename From>
-To HalfFix(From h) {
-  To ret;
-  ret.x = h.x;
-  return ret;
-}
-
-AT_API std::ostream& operator<<(std::ostream & out, const Half& value);
-
-} // namespace at
-
-#include "Half-inl.h"
-
-#undef AT_HOSTDEVICE
+#include <ATen/core/Half.h>

aten/src/ATen/InferSize.h

@@ -0,0 +1,44 @@
+#pragma once
+
+#include <c10/core/ScalarType.h>
+#include <c10/util/Optional.h>
+#include <sstream>
+#include <vector>
+
+namespace at {
+
+// Infers the size of a dim with size -1, if it exists. Also checks that new
+// shape is compatible with the number of elements.
+static std::vector<int64_t> infer_size(IntList shape, int64_t numel) {
+  auto res = shape.vec();
+  int64_t newsize = 1;
+  auto infer_dim = c10::optional<int64_t>();
+  for (int64_t dim = 0, ndim = shape.size(); dim != ndim; dim++) {
+    if (shape[dim] == -1) {
+      if (infer_dim) {
+        throw std::runtime_error("only one dimension can be inferred");
+      }
+      infer_dim = dim;
+    } else if (shape[dim] >= 0) {
+      newsize *= shape[dim];
+    } else {
+      AT_ERROR("invalid shape dimension ", shape[dim]);
+    }
+  }
+
+  if (numel == newsize || (infer_dim && newsize > 0 && numel % newsize == 0)) {
+    if (infer_dim) {
+      // we have a degree of freedom here to select the dimension size; follow NumPy semantics
+      // and just bail.
+      AT_CHECK(newsize != 0, "cannot reshape tensor of 0 elements into shape ", shape);
+      res[*infer_dim] = numel / newsize;
+    }
+    return res;
+  }
+
+  std::ostringstream ss;
+  ss << "shape '" << shape << "' is invalid for input of size " << numel;
+  throw std::runtime_error(ss.str());
+}
+
+}

aten/src/ATen/InitialTensorOptions.h

@@ -0,0 +1,15 @@
+#pragma once
+
+#include <ATen/core/TensorOptions.h>
+
+namespace at {
+
+// Represents the initial TensorOptions, before the "defaults" are ever changed.
+// This is designed to be used in library code, where the explicit devices, dtypes, etc. are known.
+// NOTE: this is not a stable API.
+inline TensorOptions initialTensorOptions() {
+  return TensorOptions(kCPU).dtype(kFloat).layout(kStrided)
+                            .requires_grad(false).is_variable(false);
+}
+
+}

aten/src/ATen/Layout.h

@@ -1,20 +1,2 @@
 #pragma once
-
-#include <ATen/ScalarType.h>
-
-namespace at {
-enum class Layout { Strided, Sparse };
-
-constexpr auto kStrided = Layout::Strided;
-constexpr auto kSparse = Layout::Sparse;
-
-inline Layout layout_from_backend(Backend backend) {
-  switch (backend) {
-    case Backend::SparseCPU:
-    case Backend::SparseCUDA:
-      return Layout::Sparse;
-    default:
-      return Layout::Strided;
-  }
-}
-} // namespace at
+#include <c10/core/Layout.h>

aten/src/ATen/LegacyTHDispatch.cpp

@@ -0,0 +1,12 @@
+#include <ATen/LegacyTHDispatch.h>
+
+namespace at {
+
+// TODO: This could be bad juju if someone calls globalContext() in the
+// destructor of an object with static lifetime.
+LegacyTHDispatch & globalLegacyTHDispatch() {
+  static LegacyTHDispatch singleton;
+  return singleton;
+}
+
+}

aten/src/ATen/LegacyTHDispatch.h

@@ -0,0 +1,91 @@
+#pragma once
+
+// LegacyTHDispatcher is the legacy mechanism for dispatching directly
+// to TH/THNN/THC/THCUNN functions in ATen, which is essentially a giant virtual
+// dispatch table for every TH function we support dynamically dispatching over.
+//
+// NB: We do not actually dispatch to *operators* here, the usual pattern is for
+// ATen operators to call this mechanism for their implementation, but the
+// operator itself is declared separately (e.g. as a native function "wrapper").
+//
+// Q: Why don't we just use LegacyTypeDispatch here?
+// A: Mainly separation of concerns:
+//   1) Type is for implementation of operators, which requires codegen of
+//      Variables, JIT, etc.  That is handled by the native function "wrappers";
+//      just calling into TH does not require that.
+//   2) Type does not require scalar-specific dispatch, whereas calling into TH
+//      does.  Thus, this separation allows us to evolve operator dispatch
+//      separately (i.e. to use the C10 dispatcher) from details of how to
+//      call TH functionality.
+//
+// The implmentation here is very similar to the LegacyTypeDispatch design, with
+// the following simplications:
+// 1) This is not required for a mobile build, so does not have to live in /core.
+// 2) Because these only contain function implementations, we do not have to
+//    handle the Variable/Tensor split; that is handled at the native function
+//    "wrapper" level.
+// 3) Because an operator must have been previously dispatched via the Type
+//    mechanism, we do need to handle device initialization.  This means it is
+//    WRONG to call directly into these functions without first going through
+//    Type dispatch (i.e. the usual operator -> Type -> LegacyTHDispatch pattern).
+// 4) Because an operator must have been previously dispatched via the Type
+//    mechanism, we do not need to handle undefined Tensors.
+//
+// NB: We don't use Registry for this, because we don't want to
+// pay for a hash table lookup every time we do an operation.
+//
+// NB: we can delete this when we don't call into any TH implementations.
+
+#include <c10/core/Backend.h>
+#include <c10/core/ScalarType.h>
+#include <ATen/LegacyTHDispatcher.h>
+
+namespace at {
+
+struct Type;
+
+struct CAFFE2_API LegacyTHDispatcherDeleter {
+  using LegacyTHDispatcherDeleterFun = void(LegacyTHDispatcher*);
+  LegacyTHDispatcherDeleterFun *fn_ = nullptr;
+  LegacyTHDispatcherDeleter() {}
+  /* implicit */ LegacyTHDispatcherDeleter(LegacyTHDispatcherDeleterFun *fn) : fn_(fn) {}
+  void operator()(LegacyTHDispatcher * ptr) {
+    if (fn_) {
+      (*fn_)(ptr);
+    }
+  }
+};
+
+class CAFFE2_API LegacyTHDispatch {
+ public:
+  using LegacyTHDispatcherUniquePtr = std::unique_ptr<LegacyTHDispatcher, LegacyTHDispatcherDeleter>;
+  // WARNING: This function has the precondition that you have
+  // initialized the type you want to call.  This initialization
+  // step is generally done by Context, or assumed because you
+  // have a Tensor and thus the Type of that Tensor must already
+  // be initialized.
+
+  void registerDispatcher(Backend b, ScalarType s, LegacyTHDispatcherUniquePtr&& t) {
+    dispatcher_registry[static_cast<int>(b)][static_cast<int>(s)] = std::move(t);
+  }
+
+  LegacyTHDispatcher* getLegacyTHDispatcherRaw(Backend p, ScalarType s) {
+    return dispatcher_registry[static_cast<int>(p)][static_cast<int>(s)].get();
+  }
+
+  LegacyTHDispatcher & getLegacyTHDispatcher(Backend p, ScalarType s) {
+    auto* type = getLegacyTHDispatcherRaw(p, s);
+    if (!type) AT_ERROR(toString(p), toString(s), "THDispatcher is not enabled.");
+    return *type;
+  }
+private:
+  // NB: dispatcher_registry has nullptr for all CUDA backends until
+  // CUDA initialization has occurred
+  LegacyTHDispatcherUniquePtr dispatcher_registry
+    [static_cast<int>(Backend::NumOptions)]
+    [static_cast<int>(ScalarType::NumOptions)];
+};
+
+CAFFE2_API LegacyTHDispatch& globalLegacyTHDispatch();
+
+} // namespace at

aten/src/ATen/MatrixRef.h

@@ -1,6 +1,6 @@
 #pragma once
-#include <ATen/ArrayRef.h>
 #include <ATen/Utils.h>
+#include <c10/util/ArrayRef.h>
 
 #include <vector>
 

aten/src/ATen/OptionsGuard.cpp

@@ -1,16 +0,0 @@
-#include <ATen/OptionsGuard.h>
-#include <ATen/optional.h>
-
-namespace at {
-
-thread_local at::optional<TensorOptions> DefaultTensorOptions::options_;
-
-TensorOptions& DefaultTensorOptions::get() {
-  if (!options_) {
-    options_.emplace(
-        /*use_thread_local_default_options=*/false);
-  }
-  return *options_;
-}
-
-} // namespace at

aten/src/ATen/OptionsGuard.h

@@ -1,54 +0,0 @@
-#pragma once
-
-#include <ATen/Device.h>
-#include <ATen/Layout.h>
-#include <ATen/ScalarType.h>
-#include <ATen/TensorOptions.h>
-#include <ATen/optional.h>
-
-namespace at {
-
-/// A wrapper over a thread local TensorOptions instance.
-struct DefaultTensorOptions {
-  /// Returns the current thread local default options.
-  /// Defined in OptionsGuard.cpp because we can't use optional in headers, due
-  /// to Windows and other compilers.
-  static TensorOptions& get();
-
- private:
-  /// This is an optional because of compiler bugs that mis-initialize static
-  /// thread local variables. The workaround is lazy initialization, i.e.
-  /// `DefaultTensorOptions::get()` will initialize the `options_` to a proper
-  /// value upon first invocation.
-  /// https://gcc.gnu.org/ml/gcc-bugs/2013-12/msg00026.html
-  static thread_local at::optional<TensorOptions> options_;
-};
-
-/// RAII guard that stores the current default options upon construction, sets
-/// the current default options to the ones given to its constructor, and
-/// finally resets the options back to the original ones in the destructor.
-struct OptionsGuard {
-  /// Stores the current default options and sets them to the given ones.
-  explicit OptionsGuard(const TensorOptions& options)
-      : original_(DefaultTensorOptions::get()) {
-    DefaultTensorOptions::get() = options;
-  }
-
-  /// Restores the original default options.
-  ~OptionsGuard() {
-    DefaultTensorOptions::get() = original_;
-  }
-
-  /// Returns the original options that were in place at the time of
-  /// construction of this object.
-  const TensorOptions& original() {
-    return original_;
-  }
-
- private:
-  /// The original options that were in place at the time of construction of
-  /// this object.
-  TensorOptions original_;
-};
-
-} // namespace at

aten/src/ATen/Parallel.h

@@ -1,6 +1,8 @@
 #pragma once
 #include <ATen/ATen.h>
+#include <atomic>
 #include <cstddef>
+#include <exception>
 
 #ifdef _OPENMP
 #include <omp.h>
@@ -20,24 +22,62 @@ inline int64_t divup(int64_t x, int64_t y) {
   return (x + y - 1) / y;
 }
 
+inline int get_max_threads() {
+#ifdef _OPENMP
+  return omp_get_max_threads();
+#else
+  return 1;
+#endif
+}
+
+inline int get_thread_num() {
+#ifdef _OPENMP
+  return omp_get_thread_num();
+#else
+  return 0;
+#endif
+}
+
+inline bool in_parallel_region() {
+#ifdef _OPENMP
+  return omp_in_parallel();
+#else
+  return false;
+#endif
+}
+
 template <class F>
 inline void parallel_for(
     const int64_t begin,
     const int64_t end,
     const int64_t grain_size,
-    const F f) {
+    const F& f) {
 #ifdef _OPENMP
-#pragma omp parallel if ((end - begin) >= grain_size)
+  std::atomic_flag err_flag = ATOMIC_FLAG_INIT;
+  std::exception_ptr eptr;
+#pragma omp parallel if (!omp_in_parallel() && ((end - begin) >= grain_size))
   {
     int64_t num_threads = omp_get_num_threads();
     int64_t tid = omp_get_thread_num();
     int64_t chunk_size = divup((end - begin), num_threads);
     int64_t begin_tid = begin + tid * chunk_size;
-    if (begin_tid < end)
-      f(begin_tid, std::min(end, chunk_size + begin_tid));
+    if (begin_tid < end) {
+      try {
+        f(begin_tid, std::min(end, chunk_size + begin_tid));
+      } catch (...) {
+        if (!err_flag.test_and_set()) {
+          eptr = std::current_exception();
+        }
+      }
+    }
+  }
+  if (eptr) {
+    std::rethrow_exception(eptr);
   }
 #else
-  f(begin, end);
+  if (begin < end) {
+    f(begin, end);
+  }
 #endif
 }
 

aten/src/ATen/Registry.h

@@ -1,216 +0,0 @@
-#pragma once
-
-/**
- * Simple registry implementation that uses static variables to
- * register object creators during program initialization time.
- */
-
-// NB: This Registry works poorly when you have other namespaces.
-// Make all macro invocations from inside the at namespace.
-
-#include <algorithm>
-#include <cstdio>
-#include <cstdlib>
-#include <functional>
-#include <memory>
-#include <mutex>
-#include <unordered_map>
-#include <string>
-#include <vector>
-
-#include <ATen/Backtrace.h>
-#include <ATen/ATenGeneral.h>
-
-namespace at {
-
-template <typename KeyType>
-inline void PrintOffendingKey(const KeyType& /*key*/) {
-  printf("[key type printing not supported]\n");
-}
-
-template <>
-inline void PrintOffendingKey(const std::string& key) {
-  printf("Offending key: %s.\n", key.c_str());
-}
-
-/**
- * @brief A template class that allows one to register classes by keys.
- *
- * The keys are usually a std::string specifying the name, but can be anything that
- * can be used in a std::map.
- *
- * You should most likely not use the Registry class explicitly, but use the
- * helper macros below to declare specific registries as well as registering
- * objects.
- */
-template <class SrcType, class ObjectPtrType, class... Args>
-class AT_API Registry {
- public:
-  typedef std::function<ObjectPtrType(Args...)> Creator;
-
-  Registry() : registry_() {}
-
-  void Register(const SrcType& key, Creator creator) {
-    // The if statement below is essentially the same as the following line:
-    // CHECK_EQ(registry_.count(key), 0) << "Key " << key
-    //                                   << " registered twice.";
-    // However, CHECK_EQ depends on google logging, and since registration is
-    // carried out at static initialization time, we do not want to have an
-    // explicit dependency on glog's initialization function.
-    std::lock_guard<std::mutex> lock(register_mutex_);
-    if (registry_.count(key) != 0) {
-      printf("Key already registered.\n");
-      PrintOffendingKey(key);
-      std::exit(1);
-    }
-    registry_[key] = creator;
-  }
-
-  void Register(const SrcType& key, Creator creator, const std::string& help_msg) {
-    Register(key, creator);
-    help_message_[key] = help_msg;
-  }
-
-  inline bool Has(const SrcType& key) { return (registry_.count(key) != 0); }
-
-  ObjectPtrType Create(const SrcType& key, Args... args) {
-    if (registry_.count(key) == 0) {
-      // Returns nullptr if the key is not registered.
-      return nullptr;
-    }
-    return registry_[key](args...);
-  }
-
-  /**
-   * Returns the keys currently registered as a std::vector.
-   */
-  std::vector<SrcType> Keys() {
-    std::vector<SrcType> keys;
-    for (const auto& it : registry_) {
-      keys.push_back(it.first);
-    }
-    return keys;
-  }
-
-  const std::unordered_map<SrcType, std::string>& HelpMessage() const {
-    return help_message_;
-  }
-
-  const char* HelpMessage(const SrcType& key) const {
-    auto it = help_message_.find(key);
-    if (it == help_message_.end()) {
-      return nullptr;
-    }
-    return it->second.c_str();
-  }
-
- private:
-  std::unordered_map<SrcType, Creator> registry_;
-  std::unordered_map<SrcType, std::string> help_message_;
-  std::mutex register_mutex_;
-
-  Registry(const Registry&) = delete;
-  Registry& operator=(const Registry&) = delete;
-};
-
-template <class SrcType, class ObjectPtrType, class... Args>
-class AT_API Registerer {
- public:
-  Registerer(
-      const SrcType& key,
-      Registry<SrcType, ObjectPtrType, Args...>* registry,
-      typename Registry<SrcType, ObjectPtrType, Args...>::Creator creator,
-      const std::string& help_msg = "") {
-    registry->Register(key, creator, help_msg);
-  }
-
-  template <class DerivedType>
-  static ObjectPtrType DefaultCreator(Args... args) {
-    // TODO(jiayq): old versions of NVCC does not handle make_unique well
-    // so we are forced to use a unique_ptr constructor here. Check if it is
-    // fine to use make_unique in the future.
-    // return make_unique<DerivedType>(args...);
-    return ObjectPtrType(new DerivedType(args...));
-  }
-};
-
-/**
- * AT_ANONYMOUS_VARIABLE(str) introduces an identifier starting with
- * str and ending with a number that varies with the line.
- * Pretty much a copy from 'folly/Preprocessor.h'
- */
-#define AT_CONCATENATE_IMPL(s1, s2) s1##s2
-#define AT_CONCATENATE(s1, s2) AT_CONCATENATE_IMPL(s1, s2)
-#ifdef __COUNTER__
-#define AT_ANONYMOUS_VARIABLE(str) AT_CONCATENATE(str, __COUNTER__)
-#else
-#define AT_ANONYMOUS_VARIABLE(str) AT_CONCATENATE(str, __LINE__)
-#endif
-
-/**
- * AT_DECLARE_TYPED_REGISTRY is a macro that expands to a function
- * declaration, as well as creating a convenient typename for its corresponding
- * registerer.
- */
-#define AT_DECLARE_TYPED_REGISTRY(                                    \
-    RegistryName, SrcType, ObjectType, PtrType, ...)                     \
-  AT_API Registry<SrcType, PtrType<ObjectType>, __VA_ARGS__>* RegistryName(); \
-  typedef Registerer<SrcType, PtrType<ObjectType>, __VA_ARGS__>        \
-      Registerer##RegistryName; \
-  extern template class Registerer<SrcType, PtrType<ObjectType>, __VA_ARGS__>;
-
-#define AT_DEFINE_TYPED_REGISTRY(                                         \
-    RegistryName, SrcType, ObjectType, PtrType, ...)                         \
-  Registry<SrcType, PtrType<ObjectType>, __VA_ARGS__>* RegistryName() {    \
-    static Registry<SrcType, PtrType<ObjectType>, __VA_ARGS__>* registry = \
-        new Registry<SrcType, PtrType<ObjectType>, __VA_ARGS__>();         \
-    return registry;                                                         \
-  } \
-  template class Registerer<SrcType, PtrType<ObjectType>, __VA_ARGS__>;
-
-// Note(Yangqing): The __VA_ARGS__ below allows one to specify a templated
-// creator with comma in its templated arguments.
-#define AT_REGISTER_TYPED_CREATOR(RegistryName, key, ...)                  \
-  namespace {                                                                 \
-  Registerer##RegistryName AT_ANONYMOUS_VARIABLE(g_##RegistryName)( \
-      key, RegistryName(), __VA_ARGS__);                                      \
-  }
-
-#define AT_REGISTER_TYPED_CLASS(RegistryName, key, ...)                    \
-  namespace {                                                                 \
-  Registerer##RegistryName AT_ANONYMOUS_VARIABLE(g_##RegistryName)( \
-      key,                                                                    \
-      RegistryName(),                                                         \
-      Registerer##RegistryName::DefaultCreator<__VA_ARGS__>,                  \
-      ::at::demangle_type<__VA_ARGS__>());                                           \
-  }
-
-// AT_DECLARE_REGISTRY and AT_DEFINE_REGISTRY are hard-wired to use std::string
-// as the key
-// type, because that is the most commonly used cases.
-#define AT_DECLARE_REGISTRY(RegistryName, ObjectType, ...) \
-  AT_DECLARE_TYPED_REGISTRY(                               \
-      RegistryName, std::string, ObjectType, std::unique_ptr, __VA_ARGS__)
-
-#define AT_DEFINE_REGISTRY(RegistryName, ObjectType, ...) \
-  AT_DEFINE_TYPED_REGISTRY(                               \
-      RegistryName, std::string, ObjectType, std::unique_ptr, __VA_ARGS__)
-
-#define AT_DECLARE_SHARED_REGISTRY(RegistryName, ObjectType, ...) \
-  AT_DECLARE_TYPED_REGISTRY(                                      \
-      RegistryName, std::string, ObjectType, std::shared_ptr, __VA_ARGS__)
-
-#define AT_DEFINE_SHARED_REGISTRY(RegistryName, ObjectType, ...) \
-  AT_DEFINE_TYPED_REGISTRY(                                      \
-      RegistryName, std::string, ObjectType, std::shared_ptr, __VA_ARGS__)
-
-// AT_REGISTER_CREATOR and AT_REGISTER_CLASS are hard-wired to use std::string
-// as the key
-// type, because that is the most commonly used cases.
-#define AT_REGISTER_CREATOR(RegistryName, key, ...) \
-  AT_REGISTER_TYPED_CREATOR(RegistryName, #key, __VA_ARGS__)
-
-#define AT_REGISTER_CLASS(RegistryName, key, ...) \
-  AT_REGISTER_TYPED_CLASS(RegistryName, #key, __VA_ARGS__)
-
-}  // namespace at

aten/src/ATen/Retainable.h

@@ -1,58 +0,0 @@
-#pragma once
-
-#include <atomic>
-
-namespace at {
-
-// base class for refcounted things, allows for collects of generic
-// refcounted objects that include tensors
-struct Retainable {
-  Retainable(): refcount(1), weak_refcount(1) {}
-  void retain() {
-    ++refcount;
-  }
-  void release() {
-    if(--refcount == 0) {
-      // If we know that this is the last reference then we can skip
-      // all the decrements and release_resources().
-      if (weak_refcount == 1) {
-        delete this;
-      } else {
-        release_resources();
-        weak_release();
-      }
-    }
-  }
-  void weak_retain() {
-    ++weak_refcount;
-  }
-  void weak_release() {
-    if (--weak_refcount == 0) {
-      delete this;
-    }
-  }
-  bool weak_lock() {
-    for (;;) {
-      auto current_refcount = refcount.load();
-      if (current_refcount == 0) return false;
-      if (refcount.compare_exchange_strong(current_refcount, current_refcount + 1)) break;
-    }
-    return true;
-  }
-  uint32_t use_count() const {
-    return refcount.load();
-  }
-  uint32_t weak_use_count() const {
-    return weak_refcount.load();
-  }
-
-  virtual void release_resources() {};
-  virtual ~Retainable() {}
-private:
-  // INVARIANT: once refcount reaches 0 it can never go up
-  // INVARIANT: weak_refcount = number of weak references + (refcount > 0 ? 1 : 0)
-  std::atomic<uint32_t> refcount;
-  std::atomic<uint32_t> weak_refcount;
-};
-
-}

aten/src/ATen/Scalar.cpp

@@ -1,21 +0,0 @@
-#include "ATen/Config.h"
-
-#include "ATen/Scalar.h"
-
-#include <TH/TH.h>
-
-#include "ATen/Tensor.h"
-#include "ATen/Context.h"
-
-namespace at {
-Tensor Scalar::toTensor() const {
-  if (Tag::HAS_t == tag) {
-    return Tensor(t);
-  } else if (Tag::HAS_d == tag) {
-    return CPU(kDouble).scalarTensor(*this);
-  } else {
-    assert(Tag::HAS_i == tag);
-    return CPU(kLong).scalarTensor(*this);
-  }
-}
-}

aten/src/ATen/Scalar.h

@@ -1,103 +1,3 @@
 #pragma once
 
-#include <assert.h>
-#include <stdint.h>
-#include <stdexcept>
-#include <string>
-#include <utility>
-
-#include "ATen/ATenGeneral.h"
-#include "ATen/Half.h"
-#include "ATen/ScalarType.h"
-#include "ATen/TensorBase.h"
-#include "ATen/Utils.h"
-
-
-namespace at {
-
-struct Tensor;
-
-class AT_API Scalar {
-public:
-  Scalar() : Scalar(int64_t(0)) {}
-
-  explicit Scalar(const detail::TensorBase & t)
-  : tag(Tag::HAS_t), t(t) {
-    AT_CHECK(t.defined(), "Attempting to create a Scalar from an undefined tensor");
-    AT_CHECK(t.dim() == 0, "Attempting to create a Scalar from a ", t.dim(), " dim tensor");
-  }
-
-#define DEFINE_IMPLICIT_CTOR(type,name,member) \
-  Scalar(type vv) \
-  : tag(Tag::HAS_##member) { \
-    v . member = convert<decltype(v.member),type>(vv); \
-  }
-
-  AT_FORALL_SCALAR_TYPES(DEFINE_IMPLICIT_CTOR)
-
-#undef DEFINE_IMPLICIT_CTOR
-
-  // return a new scalar that is guarenteed to be not backed by a tensor.
-  Scalar local() const {
-    if (Tag::HAS_t != tag) {
-      return *this;
-    }
-    return t.pImpl->localScalar();
-  }
-
-#define DEFINE_ACCESSOR(type,name,member) \
-  type to##name () const { \
-    if (Tag::HAS_t == tag) { \
-      return local().to##name(); \
-    } else if (Tag::HAS_d == tag) { \
-      return checked_convert<type, double>(v.d, #type); \
-    } else { \
-      return checked_convert<type, int64_t>(v.i, #type); \
-    } \
-  }
-
-  Tensor toTensor() const;
-
-  AT_FORALL_SCALAR_TYPES(DEFINE_ACCESSOR)
-
-  //also support scalar.to<int64_t>();
-  template<typename T>
-  T to();
-
-#undef DEFINE_ACCESSOR
-  bool isFloatingPoint() const {
-    return Tag::HAS_d == tag;
-  }
-  bool isIntegral() const {
-    return Tag::HAS_i == tag;
-  }
-  bool isBackedByTensor() const {
-    return Tag::HAS_t == tag;
-  }
-
-private:
-  enum class Tag { HAS_d, HAS_i, HAS_t };
-  Tag tag;
-  union {
-    double d;
-    int64_t i;
-  } v;
-  detail::TensorBase t;
-  friend struct Type;
-};
-
-// define the scalar.to<int64_t>() specializations
-template<typename T>
-inline T Scalar::to() {
-  throw std::runtime_error("to() cast to unexpected type.");
-}
-
-#define DEFINE_TO(T,name,_) \
-template<> \
-inline T Scalar::to<T>() { \
-  return to##name(); \
-}
-AT_FORALL_SCALAR_TYPES(DEFINE_TO)
-#undef DEFINE_TO
-
-}
+#include <ATen/core/Scalar.h>

aten/src/ATen/ScalarOps.h

@@ -0,0 +1,20 @@
+#pragma once
+
+#include <c10/core/Scalar.h>
+#include "ATen/Tensor.h"
+
+// This is in the c10 namespace because we use ADL to find the functions in it.
+namespace c10 {
+
+// FIXME: this should be (and was) Scalar::toTensor, but there is currently no way
+// to implement this without going through Derived Types (which are not part of core).
+inline at::Tensor scalar_to_tensor(Scalar s) {
+  if (s.isFloatingPoint()) {
+    return at::CPU(kDouble).scalarTensor(s);
+  } else {
+    AT_ASSERT(s.isIntegral());
+    return at::CPU(kLong).scalarTensor(s);
+  }
+}
+
+}

aten/src/ATen/ScalarType.h

@@ -1,170 +1,4 @@
 #pragma once
-
-#include <stdint.h>
-
-#include "ATen/ArrayRef.h"
-#include "ATen/ATenGeneral.h"
-#include "ATen/Half.h"
-
-namespace at {
-
-// NB: Order matters for this macro; it is relied upon in
-// _promoteTypesLookup and probably other places.
-#define AT_FORALL_SCALAR_TYPES(_) \
-_(uint8_t,Byte,i) \
-_(int8_t,Char,i) \
-_(int16_t,Short,i) \
-_(int,Int,i) \
-_(int64_t,Long,i) \
-_(at::Half,Half,d) \
-_(float,Float,d) \
-_(double,Double,d)
-
-#define AT_FORALL_SCALAR_TYPES_EXCEPT_HALF(_) \
-_(uint8_t,Byte,i) \
-_(int8_t,Char,i) \
-_(int16_t,Short,i) \
-_(int,Int,i) \
-_(int64_t,Long,i) \
-_(float,Float,d) \
-_(double,Double,d)
-
-enum class ScalarType {
-#define DEFINE_ENUM(_1,n,_2) \
-  n,
-  AT_FORALL_SCALAR_TYPES(DEFINE_ENUM)
-#undef DEFINE_ENUM
-  Undefined,
-  NumOptions
-};
-
-enum class Backend {
-  CPU,
-  CUDA,
-  SparseCPU,
-  SparseCUDA,
-  Undefined,
-  NumOptions
-};
-
-constexpr Backend kCPU = Backend::CPU;
-constexpr Backend kCUDA = Backend::CUDA;
-constexpr Backend kSparseCPU = Backend::SparseCPU;
-constexpr Backend kSparseCUDA = Backend::SparseCUDA;
-
-static inline Backend toSparse(Backend b) {
-  switch (b) {
-    case Backend::CPU: return Backend::SparseCPU;
-    case Backend::CUDA: return Backend::SparseCUDA;
-    case Backend::SparseCPU: return Backend::SparseCPU;
-    case Backend::SparseCUDA: return Backend::SparseCUDA;
-    default: throw std::runtime_error("Unknown backend");
-  }
-}
-
-static inline Backend toDense(Backend b) {
-  switch (b) {
-    case Backend::CPU: return Backend::CPU;
-    case Backend::CUDA: return Backend::CUDA;
-    case Backend::SparseCPU: return Backend::CPU;
-    case Backend::SparseCUDA: return Backend::CUDA;
-    default: throw std::runtime_error("Unknown backend");
-  }
-}
-
-static inline const char * toString(Backend b) {
-  switch(b) {
-    case Backend::CPU: return "CPU";
-    case Backend::CUDA: return "CUDA";
-    case Backend::SparseCPU: return "SparseCPU";
-    case Backend::SparseCUDA: return "SparseCUDA";
-    default: return "UNKNOWN_BACKEND";
-  }
-}
-
-#define DEFINE_CONSTANT(_,name,_2) \
-constexpr ScalarType k##name = ScalarType::name;
-
-AT_FORALL_SCALAR_TYPES(DEFINE_CONSTANT)
-#undef DEFINE_CONSTANT
-
-static inline const char * toString(ScalarType t) {
-#define DEFINE_CASE(_,name,_2) \
-  case ScalarType:: name : return #name;
-
-  switch(t) {
-    AT_FORALL_SCALAR_TYPES(DEFINE_CASE)
-    default:
-      return "UNKNOWN_SCALAR";
-  }
-#undef DEFINE_CASE
-}
-
-static inline size_t elementSize(ScalarType t) {
-#define CASE_ELEMENTSIZE_CASE(ctype,name,_2) \
-  case ScalarType:: name : return sizeof(ctype);
-
-  switch(t) {
-    AT_FORALL_SCALAR_TYPES(CASE_ELEMENTSIZE_CASE)
-    default:
-      AT_ERROR("Unknown ScalarType");
-  }
-#undef CASE_ELEMENTSIZE_CASE
-}
-
-static inline bool isIntegralType(ScalarType t) {
-  return (t == ScalarType::Byte ||
-          t == ScalarType::Char ||
-          t == ScalarType::Int ||
-          t == ScalarType::Long ||
-          t == ScalarType::Short);
-}
-
-static inline bool isFloatingType(ScalarType t) {
-  return (t == ScalarType::Double ||
-          t == ScalarType::Float ||
-          t == ScalarType::Half);
-}
-
-static inline ScalarType promoteTypes(ScalarType a, ScalarType b) {
-  // This is generated according to NumPy's promote_types
-#define u1 ScalarType::Byte
-#define i1 ScalarType::Char
-#define i2 ScalarType::Short
-#define i4 ScalarType::Int
-#define i8 ScalarType::Long
-#define f2 ScalarType::Half
-#define f4 ScalarType::Float
-#define f8 ScalarType::Double
-#define ud ScalarType::Undefined
-  static constexpr ScalarType _promoteTypesLookup
-      [static_cast<int>(ScalarType::NumOptions)]
-      [static_cast<int>(ScalarType::NumOptions)] = {
-            /* u1  i1  i2  i4  i8  f2  f4  f8, ud */
-    /* u1 */ { u1, i2, i2, i4, i8, f2, f4, f8, ud },
-    /* i1 */ { i2, i1, i2, i4, i8, f2, f4, f8, ud },
-    /* i2 */ { i2, i2, i2, i4, i8, f4, f4, f8, ud },
-    /* i4 */ { i4, i4, i4, i4, i8, f8, f4, f8, ud },
-    /* i8 */ { i8, i8, i8, i8, i8, f8, f4, f8, ud },
-    /* f2 */ { f2, f2, f4, f8, f8, f2, f4, f8, ud },
-    /* f4 */ { f4, f4, f4, f4, f4, f4, f4, f8, ud },
-    /* f8 */ { f8, f8, f8, f8, f8, f8, f8, f8, ud },
-    /* ud */ { ud, ud, ud, ud, ud, ud, ud, ud, ud },
-  };
-#undef u1
-#undef i1
-#undef i2
-#undef i4
-#undef i8
-#undef f2
-#undef f4
-#undef f8
-#undef ud
-  return _promoteTypesLookup[static_cast<int>(a)][static_cast<int>(b)];
-}
-
-struct Tensor;
-typedef ArrayRef<int64_t> IntList;
-typedef ArrayRef<Tensor> TensorList;
-
-} // namespace at
+#include <ATen/core/ATenGeneral.h> // for BC reasons
+#include <c10/core/Backend.h>
+#include <c10/core/ScalarType.h>

aten/src/ATen/SmallVector.h

@@ -1,974 +1,2 @@
-//===- llvm/ADT/SmallVector.h - 'Normally small' vectors --------*- C++ -*-===//
-//
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
-//
-//===----------------------------------------------------------------------===//
-//
-// This file defines the SmallVector class.
-//
-//===----------------------------------------------------------------------===//
-
-// ATen: modified from llvm::SmallVector.
-// replaced report_bad_alloc_error with std::bad_alloc
-// replaced isPodLike<T> with AT_IS_TRIVIALLY_COPYABLE
-// replaced iterator_range constructor with inline Container&& constructor
-// removed LLVM_NODISCARD and LLVM_ATTRIBUTE_ALWAYS_INLINE qualifiers
-// removed LLVM_UNLIKELY
-
 #pragma once
-
-#include "AlignOf.h"
-#include <algorithm>
-#include <cassert>
-#include <cstddef>
-#include <cstdlib>
-#include <cstring>
-#include <initializer_list>
-#include <iterator>
-#include <memory>
-#include <new>
-#include <type_traits>
-#include <utility>
-
-#if __GNUG__ && __GNUC__ < 5
-#define AT_IS_TRIVIALLY_COPYABLE(T) __has_trivial_copy(T)
-#else
-#define AT_IS_TRIVIALLY_COPYABLE(T) std::is_trivially_copyable<T>::value
-#endif
-
-namespace at {
-
-namespace detail {
-
-// From llvm/Support/MathExtras.h
-static inline uint64_t NextPowerOf2(uint64_t A) {
-  A |= (A >> 1);
-  A |= (A >> 2);
-  A |= (A >> 4);
-  A |= (A >> 8);
-  A |= (A >> 16);
-  A |= (A >> 32);
-  return A + 1;
-}
-
-}
-
-/// This is all the non-templated stuff common to all SmallVectors.
-class SmallVectorBase {
-protected:
-  void *BeginX, *EndX, *CapacityX;
-
-protected:
-  SmallVectorBase(void *FirstEl, size_t Size)
-    : BeginX(FirstEl), EndX(FirstEl), CapacityX((char*)FirstEl+Size) {}
-
-  /// This is an implementation of the grow() method which only works
-  /// on POD-like data types and is out of line to reduce code duplication.
-  void grow_pod(void *FirstEl, size_t MinSizeInBytes, size_t TSize);
-
-public:
-  /// This returns size()*sizeof(T).
-  size_t size_in_bytes() const {
-    return size_t((char*)EndX - (char*)BeginX);
-  }
-
-  /// capacity_in_bytes - This returns capacity()*sizeof(T).
-  size_t capacity_in_bytes() const {
-    return size_t((char*)CapacityX - (char*)BeginX);
-  }
-
-  bool empty() const { return BeginX == EndX; }
-};
-
-/// This is the part of SmallVectorTemplateBase which does not depend on whether
-/// the type T is a POD. The extra dummy template argument is used by ArrayRef
-/// to avoid unnecessarily requiring T to be complete.
-template <typename T, typename = void>
-class SmallVectorTemplateCommon : public SmallVectorBase {
-private:
-  template <typename, unsigned> friend struct SmallVectorStorage;
-
-  // Allocate raw space for N elements of type T.  If T has a ctor or dtor, we
-  // don't want it to be automatically run, so we need to represent the space as
-  // something else.  Use an array of char of sufficient alignment.
-  using U = AlignedCharArrayUnion<T>;
-  U FirstEl;
-  // Space after 'FirstEl' is clobbered, do not add any instance vars after it.
-
-protected:
-  SmallVectorTemplateCommon(size_t Size) : SmallVectorBase(&FirstEl, Size) {}
-
-  void grow_pod(size_t MinSizeInBytes, size_t TSize) {
-    SmallVectorBase::grow_pod(&FirstEl, MinSizeInBytes, TSize);
-  }
-
-  /// Return true if this is a smallvector which has not had dynamic
-  /// memory allocated for it.
-  bool isSmall() const {
-    return BeginX == static_cast<const void*>(&FirstEl);
-  }
-
-  /// Put this vector in a state of being small.
-  void resetToSmall() {
-    BeginX = EndX = CapacityX = &FirstEl;
-  }
-
-  void setEnd(T *P) { this->EndX = P; }
-
-public:
-  using size_type = size_t;
-  using difference_type = ptrdiff_t;
-  using value_type = T;
-  using iterator = T *;
-  using const_iterator = const T *;
-
-  using const_reverse_iterator = std::reverse_iterator<const_iterator>;
-  using reverse_iterator = std::reverse_iterator<iterator>;
-
-  using reference = T &;
-  using const_reference = const T &;
-  using pointer = T *;
-  using const_pointer = const T *;
-
-  // forward iterator creation methods.
-  iterator begin() { return (iterator)this->BeginX; }
-  const_iterator begin() const { return (const_iterator)this->BeginX; }
-  iterator end() { return (iterator)this->EndX; }
-  const_iterator end() const { return (const_iterator)this->EndX; }
-
-protected:
-  iterator capacity_ptr() { return (iterator)this->CapacityX; }
-  const_iterator capacity_ptr() const { return (const_iterator)this->CapacityX;}
-
-public:
-  // reverse iterator creation methods.
-  reverse_iterator rbegin()            { return reverse_iterator(end()); }
-  const_reverse_iterator rbegin() const{ return const_reverse_iterator(end()); }
-  reverse_iterator rend()              { return reverse_iterator(begin()); }
-  const_reverse_iterator rend() const { return const_reverse_iterator(begin());}
-
-  size_type size() const { return end()-begin(); }
-  size_type max_size() const { return size_type(-1) / sizeof(T); }
-
-  /// Return the total number of elements in the currently allocated buffer.
-  size_t capacity() const { return capacity_ptr() - begin(); }
-
-  /// Return a pointer to the vector's buffer, even if empty().
-  pointer data() { return pointer(begin()); }
-  /// Return a pointer to the vector's buffer, even if empty().
-  const_pointer data() const { return const_pointer(begin()); }
-
-  reference operator[](size_type idx) {
-    assert(idx < size());
-    return begin()[idx];
-  }
-  const_reference operator[](size_type idx) const {
-    assert(idx < size());
-    return begin()[idx];
-  }
-
-  reference front() {
-    assert(!empty());
-    return begin()[0];
-  }
-  const_reference front() const {
-    assert(!empty());
-    return begin()[0];
-  }
-
-  reference back() {
-    assert(!empty());
-    return end()[-1];
-  }
-  const_reference back() const {
-    assert(!empty());
-    return end()[-1];
-  }
-};
-
-/// SmallVectorTemplateBase<isPodLike = false> - This is where we put method
-/// implementations that are designed to work with non-POD-like T's.
-template <typename T, bool isPodLike>
-class SmallVectorTemplateBase : public SmallVectorTemplateCommon<T> {
-protected:
-  SmallVectorTemplateBase(size_t Size) : SmallVectorTemplateCommon<T>(Size) {}
-
-  static void destroy_range(T *S, T *E) {
-    while (S != E) {
-      --E;
-      E->~T();
-    }
-  }
-
-  /// Move the range [I, E) into the uninitialized memory starting with "Dest",
-  /// constructing elements as needed.
-  template<typename It1, typename It2>
-  static void uninitialized_move(It1 I, It1 E, It2 Dest) {
-    std::uninitialized_copy(std::make_move_iterator(I),
-                            std::make_move_iterator(E), Dest);
-  }
-
-  /// Copy the range [I, E) onto the uninitialized memory starting with "Dest",
-  /// constructing elements as needed.
-  template<typename It1, typename It2>
-  static void uninitialized_copy(It1 I, It1 E, It2 Dest) {
-    std::uninitialized_copy(I, E, Dest);
-  }
-
-  /// Grow the allocated memory (without initializing new elements), doubling
-  /// the size of the allocated memory. Guarantees space for at least one more
-  /// element, or MinSize more elements if specified.
-  void grow(size_t MinSize = 0);
-
-public:
-  void push_back(const T &Elt) {
-    if (this->EndX >= this->CapacityX)
-      this->grow();
-    ::new ((void*) this->end()) T(Elt);
-    this->setEnd(this->end()+1);
-  }
-
-  void push_back(T &&Elt) {
-    if (this->EndX >= this->CapacityX)
-      this->grow();
-    ::new ((void*) this->end()) T(::std::move(Elt));
-    this->setEnd(this->end()+1);
-  }
-
-  void pop_back() {
-    this->setEnd(this->end()-1);
-    this->end()->~T();
-  }
-};
-
-// Define this out-of-line to dissuade the C++ compiler from inlining it.
-template <typename T, bool isPodLike>
-void SmallVectorTemplateBase<T, isPodLike>::grow(size_t MinSize) {
-  size_t CurCapacity = this->capacity();
-  size_t CurSize = this->size();
-  // Always grow, even from zero.
-  size_t NewCapacity = size_t(detail::NextPowerOf2(CurCapacity+2));
-  if (NewCapacity < MinSize)
-    NewCapacity = MinSize;
-  T *NewElts = static_cast<T*>(malloc(NewCapacity*sizeof(T)));
-  if (NewElts == nullptr)
-    throw std::bad_alloc();
-
-  // Move the elements over.
-  this->uninitialized_move(this->begin(), this->end(), NewElts);
-
-  // Destroy the original elements.
-  destroy_range(this->begin(), this->end());
-
-  // If this wasn't grown from the inline copy, deallocate the old space.
-  if (!this->isSmall())
-    free(this->begin());
-
-  this->setEnd(NewElts+CurSize);
-  this->BeginX = NewElts;
-  this->CapacityX = this->begin()+NewCapacity;
-}
-
-
-/// SmallVectorTemplateBase<isPodLike = true> - This is where we put method
-/// implementations that are designed to work with POD-like T's.
-template <typename T>
-class SmallVectorTemplateBase<T, true> : public SmallVectorTemplateCommon<T> {
-protected:
-  SmallVectorTemplateBase(size_t Size) : SmallVectorTemplateCommon<T>(Size) {}
-
-  // No need to do a destroy loop for POD's.
-  static void destroy_range(T *, T *) {}
-
-  /// Move the range [I, E) onto the uninitialized memory
-  /// starting with "Dest", constructing elements into it as needed.
-  template<typename It1, typename It2>
-  static void uninitialized_move(It1 I, It1 E, It2 Dest) {
-    // Just do a copy.
-    uninitialized_copy(I, E, Dest);
-  }
-
-  /// Copy the range [I, E) onto the uninitialized memory
-  /// starting with "Dest", constructing elements into it as needed.
-  template<typename It1, typename It2>
-  static void uninitialized_copy(It1 I, It1 E, It2 Dest) {
-    // Arbitrary iterator types; just use the basic implementation.
-    std::uninitialized_copy(I, E, Dest);
-  }
-
-  /// Copy the range [I, E) onto the uninitialized memory
-  /// starting with "Dest", constructing elements into it as needed.
-  template <typename T1, typename T2>
-  static void uninitialized_copy(
-      T1 *I, T1 *E, T2 *Dest,
-      typename std::enable_if<std::is_same<typename std::remove_const<T1>::type,
-                                           T2>::value>::type * = nullptr) {
-    // Use memcpy for PODs iterated by pointers (which includes SmallVector
-    // iterators): std::uninitialized_copy optimizes to memmove, but we can
-    // use memcpy here. Note that I and E are iterators and thus might be
-    // invalid for memcpy if they are equal.
-    if (I != E)
-      memcpy(Dest, I, (E - I) * sizeof(T));
-  }
-
-  /// Double the size of the allocated memory, guaranteeing space for at
-  /// least one more element or MinSize if specified.
-  void grow(size_t MinSize = 0) {
-    this->grow_pod(MinSize*sizeof(T), sizeof(T));
-  }
-
-public:
-  void push_back(const T &Elt) {
-    if (this->EndX >= this->CapacityX)
-      this->grow();
-    memcpy(this->end(), &Elt, sizeof(T));
-    this->setEnd(this->end()+1);
-  }
-
-  void pop_back() {
-    this->setEnd(this->end()-1);
-  }
-};
-
-/// This class consists of common code factored out of the SmallVector class to
-/// reduce code duplication based on the SmallVector 'N' template parameter.
-template <typename T>
-class SmallVectorImpl : public SmallVectorTemplateBase<T, AT_IS_TRIVIALLY_COPYABLE(T)> {
-  using SuperClass = SmallVectorTemplateBase<T, AT_IS_TRIVIALLY_COPYABLE(T)>;
-
-public:
-  using iterator = typename SuperClass::iterator;
-  using const_iterator = typename SuperClass::const_iterator;
-  using size_type = typename SuperClass::size_type;
-
-protected:
-  // Default ctor - Initialize to empty.
-  explicit SmallVectorImpl(unsigned N)
-    : SmallVectorTemplateBase<T, AT_IS_TRIVIALLY_COPYABLE(T)>(N*sizeof(T)) {
-  }
-
-public:
-  SmallVectorImpl(const SmallVectorImpl &) = delete;
-
-  ~SmallVectorImpl() {
-    // Destroy the constructed elements in the vector.
-    this->destroy_range(this->begin(), this->end());
-
-    // If this wasn't grown from the inline copy, deallocate the old space.
-    if (!this->isSmall())
-      free(this->begin());
-  }
-
-  void clear() {
-    this->destroy_range(this->begin(), this->end());
-    this->EndX = this->BeginX;
-  }
-
-  void resize(size_type N) {
-    if (N < this->size()) {
-      this->destroy_range(this->begin()+N, this->end());
-      this->setEnd(this->begin()+N);
-    } else if (N > this->size()) {
-      if (this->capacity() < N)
-        this->grow(N);
-      auto I = this->end();
-      for (auto E = this->begin() + N; I != E; ++I)
-        new (&*I) T();
-      this->setEnd(this->begin()+N);
-    }
-  }
-
-  void resize(size_type N, const T &NV) {
-    if (N < this->size()) {
-      this->destroy_range(this->begin()+N, this->end());
-      this->setEnd(this->begin()+N);
-    } else if (N > this->size()) {
-      if (this->capacity() < N)
-        this->grow(N);
-      std::uninitialized_fill(this->end(), this->begin()+N, NV);
-      this->setEnd(this->begin()+N);
-    }
-  }
-
-  void reserve(size_type N) {
-    if (this->capacity() < N)
-      this->grow(N);
-  }
-
-  T pop_back_val() {
-    T Result = ::std::move(this->back());
-    this->pop_back();
-    return Result;
-  }
-
-  void swap(SmallVectorImpl &RHS);
-
-  /// Add the specified range to the end of the SmallVector.
-  template <typename in_iter,
-            typename = typename std::enable_if<std::is_convertible<
-                typename std::iterator_traits<in_iter>::iterator_category,
-                std::input_iterator_tag>::value>::type>
-  void append(in_iter in_start, in_iter in_end) {
-    size_type NumInputs = std::distance(in_start, in_end);
-    // Grow allocated space if needed.
-    if (NumInputs > size_type(this->capacity_ptr()-this->end()))
-      this->grow(this->size()+NumInputs);
-
-    // Copy the new elements over.
-    this->uninitialized_copy(in_start, in_end, this->end());
-    this->setEnd(this->end() + NumInputs);
-  }
-
-  /// Add the specified range to the end of the SmallVector.
-  void append(size_type NumInputs, const T &Elt) {
-    // Grow allocated space if needed.
-    if (NumInputs > size_type(this->capacity_ptr()-this->end()))
-      this->grow(this->size()+NumInputs);
-
-    // Copy the new elements over.
-    std::uninitialized_fill_n(this->end(), NumInputs, Elt);
-    this->setEnd(this->end() + NumInputs);
-  }
-
-  void append(std::initializer_list<T> IL) {
-    append(IL.begin(), IL.end());
-  }
-
-  // FIXME: Consider assigning over existing elements, rather than clearing &
-  // re-initializing them - for all assign(...) variants.
-
-  void assign(size_type NumElts, const T &Elt) {
-    clear();
-    if (this->capacity() < NumElts)
-      this->grow(NumElts);
-    this->setEnd(this->begin()+NumElts);
-    std::uninitialized_fill(this->begin(), this->end(), Elt);
-  }
-
-  template <typename in_iter,
-            typename = typename std::enable_if<std::is_convertible<
-                typename std::iterator_traits<in_iter>::iterator_category,
-                std::input_iterator_tag>::value>::type>
-  void assign(in_iter in_start, in_iter in_end) {
-    clear();
-    append(in_start, in_end);
-  }
-
-  void assign(std::initializer_list<T> IL) {
-    clear();
-    append(IL);
-  }
-
-  iterator erase(const_iterator CI) {
-    // Just cast away constness because this is a non-const member function.
-    iterator I = const_cast<iterator>(CI);
-
-    assert(I >= this->begin() && "Iterator to erase is out of bounds.");
-    assert(I < this->end() && "Erasing at past-the-end iterator.");
-
-    iterator N = I;
-    // Shift all elts down one.
-    std::move(I+1, this->end(), I);
-    // Drop the last elt.
-    this->pop_back();
-    return(N);
-  }
-
-  iterator erase(const_iterator CS, const_iterator CE) {
-    // Just cast away constness because this is a non-const member function.
-    iterator S = const_cast<iterator>(CS);
-    iterator E = const_cast<iterator>(CE);
-
-    assert(S >= this->begin() && "Range to erase is out of bounds.");
-    assert(S <= E && "Trying to erase invalid range.");
-    assert(E <= this->end() && "Trying to erase past the end.");
-
-    iterator N = S;
-    // Shift all elts down.
-    iterator I = std::move(E, this->end(), S);
-    // Drop the last elts.
-    this->destroy_range(I, this->end());
-    this->setEnd(I);
-    return(N);
-  }
-
-  iterator insert(iterator I, T &&Elt) {
-    if (I == this->end()) {  // Important special case for empty vector.
-      this->push_back(::std::move(Elt));
-      return this->end()-1;
-    }
-
-    assert(I >= this->begin() && "Insertion iterator is out of bounds.");
-    assert(I <= this->end() && "Inserting past the end of the vector.");
-
-    if (this->EndX >= this->CapacityX) {
-      size_t EltNo = I-this->begin();
-      this->grow();
-      I = this->begin()+EltNo;
-    }
-
-    ::new ((void*) this->end()) T(::std::move(this->back()));
-    // Push everything else over.
-    std::move_backward(I, this->end()-1, this->end());
-    this->setEnd(this->end()+1);
-
-    // If we just moved the element we're inserting, be sure to update
-    // the reference.
-    T *EltPtr = &Elt;
-    if (I <= EltPtr && EltPtr < this->EndX)
-      ++EltPtr;
-
-    *I = ::std::move(*EltPtr);
-    return I;
-  }
-
-  iterator insert(iterator I, const T &Elt) {
-    if (I == this->end()) {  // Important special case for empty vector.
-      this->push_back(Elt);
-      return this->end()-1;
-    }
-
-    assert(I >= this->begin() && "Insertion iterator is out of bounds.");
-    assert(I <= this->end() && "Inserting past the end of the vector.");
-
-    if (this->EndX >= this->CapacityX) {
-      size_t EltNo = I-this->begin();
-      this->grow();
-      I = this->begin()+EltNo;
-    }
-    ::new ((void*) this->end()) T(std::move(this->back()));
-    // Push everything else over.
-    std::move_backward(I, this->end()-1, this->end());
-    this->setEnd(this->end()+1);
-
-    // If we just moved the element we're inserting, be sure to update
-    // the reference.
-    const T *EltPtr = &Elt;
-    if (I <= EltPtr && EltPtr < this->EndX)
-      ++EltPtr;
-
-    *I = *EltPtr;
-    return I;
-  }
-
-  iterator insert(iterator I, size_type NumToInsert, const T &Elt) {
-    // Convert iterator to elt# to avoid invalidating iterator when we reserve()
-    size_t InsertElt = I - this->begin();
-
-    if (I == this->end()) {  // Important special case for empty vector.
-      append(NumToInsert, Elt);
-      return this->begin()+InsertElt;
-    }
-
-    assert(I >= this->begin() && "Insertion iterator is out of bounds.");
-    assert(I <= this->end() && "Inserting past the end of the vector.");
-
-    // Ensure there is enough space.
-    reserve(this->size() + NumToInsert);
-
-    // Uninvalidate the iterator.
-    I = this->begin()+InsertElt;
-
-    // If there are more elements between the insertion point and the end of the
-    // range than there are being inserted, we can use a simple approach to
-    // insertion.  Since we already reserved space, we know that this won't
-    // reallocate the vector.
-    if (size_t(this->end()-I) >= NumToInsert) {
-      T *OldEnd = this->end();
-      append(std::move_iterator<iterator>(this->end() - NumToInsert),
-             std::move_iterator<iterator>(this->end()));
-
-      // Copy the existing elements that get replaced.
-      std::move_backward(I, OldEnd-NumToInsert, OldEnd);
-
-      std::fill_n(I, NumToInsert, Elt);
-      return I;
-    }
-
-    // Otherwise, we're inserting more elements than exist already, and we're
-    // not inserting at the end.
-
-    // Move over the elements that we're about to overwrite.
-    T *OldEnd = this->end();
-    this->setEnd(this->end() + NumToInsert);
-    size_t NumOverwritten = OldEnd-I;
-    this->uninitialized_move(I, OldEnd, this->end()-NumOverwritten);
-
-    // Replace the overwritten part.
-    std::fill_n(I, NumOverwritten, Elt);
-
-    // Insert the non-overwritten middle part.
-    std::uninitialized_fill_n(OldEnd, NumToInsert-NumOverwritten, Elt);
-    return I;
-  }
-
-  template <typename ItTy,
-            typename = typename std::enable_if<std::is_convertible<
-                typename std::iterator_traits<ItTy>::iterator_category,
-                std::input_iterator_tag>::value>::type>
-  iterator insert(iterator I, ItTy From, ItTy To) {
-    // Convert iterator to elt# to avoid invalidating iterator when we reserve()
-    size_t InsertElt = I - this->begin();
-
-    if (I == this->end()) {  // Important special case for empty vector.
-      append(From, To);
-      return this->begin()+InsertElt;
-    }
-
-    assert(I >= this->begin() && "Insertion iterator is out of bounds.");
-    assert(I <= this->end() && "Inserting past the end of the vector.");
-
-    size_t NumToInsert = std::distance(From, To);
-
-    // Ensure there is enough space.
-    reserve(this->size() + NumToInsert);
-
-    // Uninvalidate the iterator.
-    I = this->begin()+InsertElt;
-
-    // If there are more elements between the insertion point and the end of the
-    // range than there are being inserted, we can use a simple approach to
-    // insertion.  Since we already reserved space, we know that this won't
-    // reallocate the vector.
-    if (size_t(this->end()-I) >= NumToInsert) {
-      T *OldEnd = this->end();
-      append(std::move_iterator<iterator>(this->end() - NumToInsert),
-             std::move_iterator<iterator>(this->end()));
-
-      // Copy the existing elements that get replaced.
-      std::move_backward(I, OldEnd-NumToInsert, OldEnd);
-
-      std::copy(From, To, I);
-      return I;
-    }
-
-    // Otherwise, we're inserting more elements than exist already, and we're
-    // not inserting at the end.
-
-    // Move over the elements that we're about to overwrite.
-    T *OldEnd = this->end();
-    this->setEnd(this->end() + NumToInsert);
-    size_t NumOverwritten = OldEnd-I;
-    this->uninitialized_move(I, OldEnd, this->end()-NumOverwritten);
-
-    // Replace the overwritten part.
-    for (T *J = I; NumOverwritten > 0; --NumOverwritten) {
-      *J = *From;
-      ++J; ++From;
-    }
-
-    // Insert the non-overwritten middle part.
-    this->uninitialized_copy(From, To, OldEnd);
-    return I;
-  }
-
-  void insert(iterator I, std::initializer_list<T> IL) {
-    insert(I, IL.begin(), IL.end());
-  }
-
-  template <typename... ArgTypes> void emplace_back(ArgTypes &&... Args) {
-    if (this->EndX >= this->CapacityX)
-      this->grow();
-    ::new ((void *)this->end()) T(std::forward<ArgTypes>(Args)...);
-    this->setEnd(this->end() + 1);
-  }
-
-  SmallVectorImpl &operator=(const SmallVectorImpl &RHS);
-
-  SmallVectorImpl &operator=(SmallVectorImpl &&RHS);
-
-  bool operator==(const SmallVectorImpl &RHS) const {
-    if (this->size() != RHS.size()) return false;
-    return std::equal(this->begin(), this->end(), RHS.begin());
-  }
-  bool operator!=(const SmallVectorImpl &RHS) const {
-    return !(*this == RHS);
-  }
-
-  bool operator<(const SmallVectorImpl &RHS) const {
-    return std::lexicographical_compare(this->begin(), this->end(),
-                                        RHS.begin(), RHS.end());
-  }
-
-  /// Set the array size to \p N, which the current array must have enough
-  /// capacity for.
-  ///
-  /// This does not construct or destroy any elements in the vector.
-  ///
-  /// Clients can use this in conjunction with capacity() to write past the end
-  /// of the buffer when they know that more elements are available, and only
-  /// update the size later. This avoids the cost of value initializing elements
-  /// which will only be overwritten.
-  void set_size(size_type N) {
-    assert(N <= this->capacity());
-    this->setEnd(this->begin() + N);
-  }
-};
-
-template <typename T>
-void SmallVectorImpl<T>::swap(SmallVectorImpl<T> &RHS) {
-  if (this == &RHS) return;
-
-  // We can only avoid copying elements if neither vector is small.
-  if (!this->isSmall() && !RHS.isSmall()) {
-    std::swap(this->BeginX, RHS.BeginX);
-    std::swap(this->EndX, RHS.EndX);
-    std::swap(this->CapacityX, RHS.CapacityX);
-    return;
-  }
-  if (RHS.size() > this->capacity())
-    this->grow(RHS.size());
-  if (this->size() > RHS.capacity())
-    RHS.grow(this->size());
-
-  // Swap the shared elements.
-  size_t NumShared = this->size();
-  if (NumShared > RHS.size()) NumShared = RHS.size();
-  for (size_type i = 0; i != NumShared; ++i)
-    std::swap((*this)[i], RHS[i]);
-
-  // Copy over the extra elts.
-  if (this->size() > RHS.size()) {
-    size_t EltDiff = this->size() - RHS.size();
-    this->uninitialized_copy(this->begin()+NumShared, this->end(), RHS.end());
-    RHS.setEnd(RHS.end()+EltDiff);
-    this->destroy_range(this->begin()+NumShared, this->end());
-    this->setEnd(this->begin()+NumShared);
-  } else if (RHS.size() > this->size()) {
-    size_t EltDiff = RHS.size() - this->size();
-    this->uninitialized_copy(RHS.begin()+NumShared, RHS.end(), this->end());
-    this->setEnd(this->end() + EltDiff);
-    this->destroy_range(RHS.begin()+NumShared, RHS.end());
-    RHS.setEnd(RHS.begin()+NumShared);
-  }
-}
-
-template <typename T>
-SmallVectorImpl<T> &SmallVectorImpl<T>::
-  operator=(const SmallVectorImpl<T> &RHS) {
-  // Avoid self-assignment.
-  if (this == &RHS) return *this;
-
-  // If we already have sufficient space, assign the common elements, then
-  // destroy any excess.
-  size_t RHSSize = RHS.size();
-  size_t CurSize = this->size();
-  if (CurSize >= RHSSize) {
-    // Assign common elements.
-    iterator NewEnd;
-    if (RHSSize)
-      NewEnd = std::copy(RHS.begin(), RHS.begin()+RHSSize, this->begin());
-    else
-      NewEnd = this->begin();
-
-    // Destroy excess elements.
-    this->destroy_range(NewEnd, this->end());
-
-    // Trim.
-    this->setEnd(NewEnd);
-    return *this;
-  }
-
-  // If we have to grow to have enough elements, destroy the current elements.
-  // This allows us to avoid copying them during the grow.
-  // FIXME: don't do this if they're efficiently moveable.
-  if (this->capacity() < RHSSize) {
-    // Destroy current elements.
-    this->destroy_range(this->begin(), this->end());
-    this->setEnd(this->begin());
-    CurSize = 0;
-    this->grow(RHSSize);
-  } else if (CurSize) {
-    // Otherwise, use assignment for the already-constructed elements.
-    std::copy(RHS.begin(), RHS.begin()+CurSize, this->begin());
-  }
-
-  // Copy construct the new elements in place.
-  this->uninitialized_copy(RHS.begin()+CurSize, RHS.end(),
-                           this->begin()+CurSize);
-
-  // Set end.
-  this->setEnd(this->begin()+RHSSize);
-  return *this;
-}
-
-template <typename T>
-SmallVectorImpl<T> &SmallVectorImpl<T>::operator=(SmallVectorImpl<T> &&RHS) {
-  // Avoid self-assignment.
-  if (this == &RHS) return *this;
-
-  // If the RHS isn't small, clear this vector and then steal its buffer.
-  if (!RHS.isSmall()) {
-    this->destroy_range(this->begin(), this->end());
-    if (!this->isSmall()) free(this->begin());
-    this->BeginX = RHS.BeginX;
-    this->EndX = RHS.EndX;
-    this->CapacityX = RHS.CapacityX;
-    RHS.resetToSmall();
-    return *this;
-  }
-
-  // If we already have sufficient space, assign the common elements, then
-  // destroy any excess.
-  size_t RHSSize = RHS.size();
-  size_t CurSize = this->size();
-  if (CurSize >= RHSSize) {
-    // Assign common elements.
-    iterator NewEnd = this->begin();
-    if (RHSSize)
-      NewEnd = std::move(RHS.begin(), RHS.end(), NewEnd);
-
-    // Destroy excess elements and trim the bounds.
-    this->destroy_range(NewEnd, this->end());
-    this->setEnd(NewEnd);
-
-    // Clear the RHS.
-    RHS.clear();
-
-    return *this;
-  }
-
-  // If we have to grow to have enough elements, destroy the current elements.
-  // This allows us to avoid copying them during the grow.
-  // FIXME: this may not actually make any sense if we can efficiently move
-  // elements.
-  if (this->capacity() < RHSSize) {
-    // Destroy current elements.
-    this->destroy_range(this->begin(), this->end());
-    this->setEnd(this->begin());
-    CurSize = 0;
-    this->grow(RHSSize);
-  } else if (CurSize) {
-    // Otherwise, use assignment for the already-constructed elements.
-    std::move(RHS.begin(), RHS.begin()+CurSize, this->begin());
-  }
-
-  // Move-construct the new elements in place.
-  this->uninitialized_move(RHS.begin()+CurSize, RHS.end(),
-                           this->begin()+CurSize);
-
-  // Set end.
-  this->setEnd(this->begin()+RHSSize);
-
-  RHS.clear();
-  return *this;
-}
-
-/// Storage for the SmallVector elements which aren't contained in
-/// SmallVectorTemplateCommon. There are 'N-1' elements here. The remaining '1'
-/// element is in the base class. This is specialized for the N=1 and N=0 cases
-/// to avoid allocating unnecessary storage.
-template <typename T, unsigned N>
-struct SmallVectorStorage {
-  typename SmallVectorTemplateCommon<T>::U InlineElts[N - 1];
-};
-template <typename T> struct SmallVectorStorage<T, 1> {};
-template <typename T> struct SmallVectorStorage<T, 0> {};
-
-/// This is a 'vector' (really, a variable-sized array), optimized
-/// for the case when the array is small.  It contains some number of elements
-/// in-place, which allows it to avoid heap allocation when the actual number of
-/// elements is below that threshold.  This allows normal "small" cases to be
-/// fast without losing generality for large inputs.
-///
-/// Note that this does not attempt to be exception safe.
-///
-template <typename T, unsigned N>
-class SmallVector : public SmallVectorImpl<T> {
-  /// Inline space for elements which aren't stored in the base class.
-  SmallVectorStorage<T, N> Storage;
-
-public:
-  SmallVector() : SmallVectorImpl<T>(N) {}
-
-  explicit SmallVector(size_t Size, const T &Value = T())
-    : SmallVectorImpl<T>(N) {
-    this->assign(Size, Value);
-  }
-
-  template <typename ItTy,
-            typename = typename std::enable_if<std::is_convertible<
-                typename std::iterator_traits<ItTy>::iterator_category,
-                std::input_iterator_tag>::value>::type>
-  SmallVector(ItTy S, ItTy E) : SmallVectorImpl<T>(N) {
-    this->append(S, E);
-  }
-
-  template <typename Container>
-  explicit SmallVector(Container &&c) : SmallVectorImpl<T>(N) {
-    this->append(c.begin(), c.end());
-  }
-
-  SmallVector(std::initializer_list<T> IL) : SmallVectorImpl<T>(N) {
-    this->assign(IL);
-  }
-
-  SmallVector(const SmallVector &RHS) : SmallVectorImpl<T>(N) {
-    if (!RHS.empty())
-      SmallVectorImpl<T>::operator=(RHS);
-  }
-
-  const SmallVector &operator=(const SmallVector &RHS) {
-    SmallVectorImpl<T>::operator=(RHS);
-    return *this;
-  }
-
-  SmallVector(SmallVector &&RHS) : SmallVectorImpl<T>(N) {
-    if (!RHS.empty())
-      SmallVectorImpl<T>::operator=(::std::move(RHS));
-  }
-
-  template<typename Container>
-  const SmallVector &operator=(const Container &RHS) {
-    this->assign(RHS.begin(), RHS.end());
-    return *this;
-  }
-
-  SmallVector(SmallVectorImpl<T> &&RHS) : SmallVectorImpl<T>(N) {
-    if (!RHS.empty())
-      SmallVectorImpl<T>::operator=(::std::move(RHS));
-  }
-
-  const SmallVector &operator=(SmallVector &&RHS) {
-    SmallVectorImpl<T>::operator=(::std::move(RHS));
-    return *this;
-  }
-
-  const SmallVector &operator=(SmallVectorImpl<T> &&RHS) {
-    SmallVectorImpl<T>::operator=(::std::move(RHS));
-    return *this;
-  }
-
-  const SmallVector &operator=(std::initializer_list<T> IL) {
-    this->assign(IL);
-    return *this;
-  }
-};
-
-template <typename T, unsigned N>
-inline size_t capacity_in_bytes(const SmallVector<T, N> &X) {
-  return X.capacity_in_bytes();
-}
-
-} // end namespace at
-
-namespace std {
-
-  /// Implement std::swap in terms of SmallVector swap.
-  template<typename T>
-  inline void
-  swap(at::SmallVectorImpl<T> &LHS, at::SmallVectorImpl<T> &RHS) {
-    LHS.swap(RHS);
-  }
-
-  /// Implement std::swap in terms of SmallVector swap.
-  template<typename T, unsigned N>
-  inline void
-  swap(at::SmallVector<T, N> &LHS, at::SmallVector<T, N> &RHS) {
-    LHS.swap(RHS);
-  }
-
-} // end namespace std
+#include <c10/util/SmallVector.h>

aten/src/ATen/SparseTensorImpl.cpp

@@ -1,85 +1,117 @@
 #include <ATen/ATen.h>
 #include <ATen/SparseTensorImpl.h>
+#include <ATen/InitialTensorOptions.h>
+#include <ATen/core/LegacyTypeDispatch.h>
 
 namespace at {
 
+namespace {
+  DeviceType sparseTensorIdToDeviceType(TensorTypeId type_id) {
+    if (type_id == SparseCPUTensorId()) {
+      return kCPU;
+    } else if (type_id == SparseCUDATensorId()) {
+      return kCUDA;
+    } else {
+      AT_ERROR("Cannot construct SparseTensor with non-sparse tensor type ID ", type_id);
+    }
+  }
+}
+
 
 // An empty dense tensor defaults to a 1-dimensional tensor of size [0]
 // (recall, it is not a 0-dimensional tensor, because such a tensor would
 // a scalar and have one element)
 //
 // Thus, an empty sparse tensor should be a 1-dimensional tensor of size [0].
-// Furthermore, we have dim == sparseDims + denseDims; since this is a sparse
-// tensor, let us say that an empty sparse tensor has sparseDims == 1 and
-// denseDims == 0.  (There is a degree of freedom here, but given that this
-// is a sparse dimension, it seems reasonable to demand that sparseDims > 0).
+// Furthermore, we have dim == sparse_dim + dense_dim; since this is a sparse
+// tensor, let us say that an empty sparse tensor has sparse_dim == 1 and
+// dense_dim == 0.  (There is a degree of freedom here, but given that this
+// is a sparse dimension, it seems reasonable to demand that sparse_dim > 0).
 //
-// In an ideal world, this would then mean we allocate a [1,0] size indices
-// tensor and a [0] size values tensor for such an empty tensor.  However,
-// 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(Type * type)
-    : TensorImpl(type)
+// This means that we allocate a [1,0] size indices tensor and a [0] size
+// values tensor for such an empty tensor.
+SparseTensorImpl::SparseTensorImpl(at::TensorTypeId type_id, const caffe2::TypeMeta& data_type)
+    : TensorImpl(type_id, data_type, nullptr, false)
     , size_{0}
-    , sparseDims_(1)
-    , denseDims_(0)
-    , indices_(type->toDense().toScalarType(ScalarType::Long).tensor())
-    , values_(type->toDense().tensor()) {
-      AT_ASSERT(type->is_sparse());
-    }
+    , sparse_dim_(1)
+    , dense_dim_(0)
+    , indices_(at::empty({1, 0}, at::initialTensorOptions().device(sparseTensorIdToDeviceType(type_id)).dtype(ScalarType::Long)))
+    , values_(at::empty({0}, at::initialTensorOptions().device(sparseTensorIdToDeviceType(type_id)).dtype(data_type))) {}
 
-const char * SparseTensorImpl::toString() const {
-  // TODO: also give back type information
-  return "SparseTensor";
-}
 IntList SparseTensorImpl::sizes() const {
   return size_;
 }
 IntList SparseTensorImpl::strides() const {
   AT_ERROR("sparse tensors do not have strides");
 }
-int64_t SparseTensorImpl::dim() const {
-  return sparseDims_ + denseDims_;
+bool SparseTensorImpl::is_contiguous() const {
+  AT_ERROR("sparse tensors do not have is_contiguous");
 }
-Scalar SparseTensorImpl::localScalar() {
-  int64_t n = numel();
-  AT_CHECK(n == 1, "a Tensor with ", n, " elements cannot be converted to Scalar");
-  if (nnz_ == 0) return Scalar(0);
-  if (coalesced_) return values_.pImpl->localScalar();
-  // You have a non-coalesced scalar sparse tensor?!  Wow!  Have
-  // a cookie.
-  return values_.sum().pImpl->localScalar();
+int64_t SparseTensorImpl::size(int64_t d) const {
+  d = at::maybe_wrap_dim(d, dim(), false);
+  return size_[d];
 }
-void * SparseTensorImpl::unsafeGetTH(bool retain) {
-  AT_ERROR("unsafeGetTH not supported for new style TensorImpl");
+int64_t SparseTensorImpl::stride(int64_t d) const {
+  AT_ERROR("sparse tensors do not have strides");
 }
-std::unique_ptr<Storage> SparseTensorImpl::storage() {
-  AT_ERROR("sparse tensors do not have storage");
+void SparseTensorImpl::resize_dim(int64_t ndim) {
+  AT_ERROR("sparse tensors do not have resize_dim");
+}
+void SparseTensorImpl::set_size(int64_t dim, int64_t new_size) {
+  AT_ERROR("sparse tensors do not have set_size");
+}
+void SparseTensorImpl::set_stride(int64_t dim, int64_t new_stride) {
+  AT_ERROR("sparse tensors do not have set_stride");
+}
+void SparseTensorImpl::set_storage_offset(int64_t storage_offset) {
+  AT_ERROR("sparse tensors do not have set_storage_offset");
 }
 
-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
-  bool empty = values.numel() == 0;
-  AT_CHECK(values.type().toSparse() == type(), "values type must match sparse tensor type");
+int64_t SparseTensorImpl::dim() const {
+  return sparse_dim_ + dense_dim_;
+}
+TensorImpl* SparseTensorImpl::maybe_zero_dim(bool condition_when_zero_dim) {
+  AT_CHECK(condition_when_zero_dim == (dim() == 0),
+           "Attempted to maybe_zero_dim on a SparseTensorImpl to ", condition_when_zero_dim,
+           " but the SparseTensor's dim() is ", dim(), " and SparseTensors do not support"
+           " changing dimensionality via maybe_zero_dim");
+  return this;
+}
+const Storage& SparseTensorImpl::storage() 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_unsafe(const Tensor& indices, const Tensor& values) {
+  AT_ASSERT(!indices.is_variable() && !values.is_variable());  // They should be plain tensors!
+
+  AT_CHECK(!indices.is_sparse(), "expected indices to be a dense tensor, but got indices of layout ", indices.layout());
+  AT_CHECK(!values.is_sparse(), "expected values to be a dense tensor, but got values of layout ", values.layout());
+
+  AT_CHECK(values.type().toSparse() == legacyTensorType(*this), "values type must match sparse tensor type");
   AT_CHECK(indices.type().scalarType() == kLong, "indices must be an int64 tensor");
   AT_CHECK(indices.type().backend() == values.type().backend(), "backend of indices (", indices.type().backend(), ") must match backend of values (", values.type().backend(), ")");
   AT_CHECK(!indices.is_cuda() || indices.get_device() == values.get_device(), "device of indices (", indices.get_device(), ") must match device of values (", values.get_device(), ")");
-  if (!empty) {
-    AT_CHECK(indices.dim() == 2, "indices must be nDim x nnz");
-    AT_CHECK(indices.size(1) == values.size(0), "indices and values must have same nnz");
-    AT_CHECK(indices.size(0) == sparseDims_, "indices has incorrect first dimension, expected ", sparseDims_, ", got ", indices.size(0));
-    AT_CHECK(values.dim() == denseDims_ + 1, "values has incorrect number of dimensions, expected ", denseDims_ + 1, ", got ", values.dim());
-  } else {
-    AT_CHECK(indices.numel() == 0, "if values is empty, indices must be empty too");
-  }
+
+  AT_CHECK(indices.dim() == 2, "indices must be sparse_dim x nnz, but got: ", indices.sizes());
+  AT_CHECK(indices.size(1) == values.size(0), "indices and values must have same nnz, but got nnz from indices: ", indices.size(1), ", nnz from values: ", values.size(0));
+  AT_CHECK(indices.size(0) == sparse_dim_, "indices has incorrect first dimension, expected ", sparse_dim_, ", got ", indices.size(0));
+  AT_CHECK(values.dim() == dense_dim_ + 1, "values has incorrect number of dimensions, expected ", dense_dim_ + 1, ", got ", values.dim());
+
+  auto dense_size_original = sizes().slice(sparse_dim_);
+  std::vector<int64_t> expected_values_size_vec = {values.size(0)};
+  expected_values_size_vec.insert(expected_values_size_vec.end(), dense_size_original.begin(), dense_size_original.end());
+  IntList expected_values_size(expected_values_size_vec);
+  auto new_values_size = values.sizes();
+  AT_CHECK(
+    std::equal(expected_values_size.begin(), expected_values_size.end(), new_values_size.begin()),
+    "values has incorrect size, expected ", expected_values_size, ", got ", new_values_size
+  );
+
   indices_ = indices;
   values_ = values;
-  // TODO: Eliminate this ternary when we handle size zero dimensions.
-  // (Actually, this will "accidentally" work today because all zero-size
-  // tensors have size [0], and so you'll get 0 when empty is zero; but it's
-  // more explicit this way.)
-  nnz_ = empty ? 0 : values.size(0);
+
   coalesced_ = false;
 }
 

aten/src/ATen/SparseTensorImpl.h

@@ -1,38 +1,26 @@
 #pragma once
 
 #include "ATen/Tensor.h"
-#include "ATen/TensorImpl.h"
-#include "ATen/Error.h"
+#include "ATen/core/TensorImpl.h"
+#include "c10/util/Exception.h"
 
 namespace at {
-struct SparseTensorImpl : public TensorImpl {
+struct CAFFE2_API SparseTensorImpl : public TensorImpl {
   // Stored in COO format, indices + values.
 
-  // Ideal INVARIANTS:
-  // _sparseDims: range [0, len(shape)]; _sparseDims + _denseDims = len(shape)
-  // _denseDims : range [0, len(shape)]; _sparseDims + _denseDims = len(shape)
-  // _indices.shape: dimensionality: 2,  shape: (_sparseDims, nnz)
-  // _values.shape:  dimensionality: 1 + _denseDims.  shape: (nnz, shape[_sparseDims:])
-
-  // Actual INVARIANT differences:
-  // 1) _sparseDims: range [1, len(shape)] (i.e. we don't allow 0 sparse dimensions)
-  // 2) when nnz = 0, there is strange behavior because we lack 0-dimensional sparse tensors.  Namely:
-  //    dimensionality == 0, _sparseDims == 0, _denseDims == 0, _indices.shape == {0}, _values.shape == {0}
-  // 3) For both _indices.shape and _values.shape, the nnz dimension may be larger than nnz
-  // 4) For _values.shape, the non-nnz dimensions may be smaller than the corresponding dimension size, e.g.
-  //    a shape (2,3) sparse tensor with _sparseDims == 1, may have _values.shape: (nnz, <=2, <=3).
-
+  // INVARIANTS:
+  // sparse_dim: range [0, len(shape)]; sparse_dim + dense_dim = len(shape)
+  // dense_dim : range [0, len(shape)]; sparse_dim + dense_dim = len(shape)
+  // _indices.shape: dimensionality: 2,  shape: (sparse_dim, nnz)
+  // _values.shape:  dimensionality: 1 + dense_dim.  shape: (nnz, shape[sparse_dim:])
 
   // The true size of the sparse tensor (e.g., if you called to_dense()
   // on it).  When THTensor merges into TensorImpl, this field
   // should move to the parent class.
   std::vector<int64_t> size_;
 
-  // The number of non-zero elements.
-  int64_t nnz_ = 0;
-
-  int64_t sparseDims_ = 0; // number of sparse dimensions
-  int64_t denseDims_ = 0; // number of dense dimensions
+  int64_t sparse_dim_ = 0; // number of sparse dimensions
+  int64_t dense_dim_ = 0; // number of dense dimensions
 
   Tensor indices_; // always a LongTensor
   Tensor values_;
@@ -48,58 +36,157 @@ struct SparseTensorImpl : public TensorImpl {
 
 public:
   // Public for now...
-  explicit SparseTensorImpl(Type * type);
+  explicit SparseTensorImpl(at::TensorTypeId, const caffe2::TypeMeta&);
 
-  int64_t nnz() const { return nnz_; }
-  int64_t sparseDims() const { return sparseDims_; }
-  int64_t denseDims() const { return denseDims_; }
+  int64_t nnz() const { return values_.size(0); }
+  int64_t sparse_dim() const { return sparse_dim_; }
+  int64_t dense_dim() const { return dense_dim_; }
   bool coalesced() const { return coalesced_; }
   Tensor indices() const { return indices_; }
   Tensor values() const { return values_; }
 
-  const char * toString() const override;
   IntList sizes() const override;
   IntList strides() const override;
+  bool is_contiguous() const override;
+  int64_t size(int64_t d) const override;
+  int64_t stride(int64_t d) const override;
+  void resize_dim(int64_t ndim) override;
+  void set_size(int64_t dim, int64_t new_size) override;
+  void set_stride(int64_t dim, int64_t new_stride) override;
+  void set_storage_offset(int64_t storage_offset) override;
+
   int64_t dim() const override;
-  Scalar localScalar() override;
-  void * unsafeGetTH(bool retain) override;
-  std::unique_ptr<Storage> storage() override;
+  TensorImpl* maybe_zero_dim(bool condition_when_zero_dim) override;
+  const Storage& storage() 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
-  std::vector<int64_t>& _sizes_mut() { return size_; }
-
-  // WARNING: This function does NOT preserve invariants of sparseDims/denseDims with
+  // WARNING: This function does NOT preserve invariants of sparse_dim/dense_dim with
   // respect to indices and values
-  void raw_resize_(int64_t sparseDims, int64_t denseDims, ArrayRef<int64_t> size) {
-    // UGHHHHH.  Legacy special case
-    if (size.size() == 0) {
-      size_ = {0};
-    } else {
-      size_ = size;
-    }
-    sparseDims_ = sparseDims;
-    denseDims_ = denseDims;
+  void raw_resize_(int64_t sparse_dim, int64_t dense_dim, IntList size) {
+    size_ = size.vec();
+    sparse_dim_ = sparse_dim;
+    dense_dim_ = dense_dim;
+    refresh_numel();
   }
 
-  // TODO: I hate these two setters, please get rid of them!!!
-  void set_indices(const Tensor& indices) {
-    AT_ASSERT(indices.type().backend() == at::toDense(type().backend()));
-    AT_ASSERT(indices.type().scalarType() == kLong);
-    indices_ = indices;
+  // NOTE: This function preserves invariants of sparse_dim/dense_dim with respect to
+  // indices and values.
+  //
+  // NOTE: This function supports the following cases:
+  // 1. When we keep the number of dense dimensions unchanged, and NOT shrinking the size of
+  // any of the dense dimensions.
+  // 2. When we keep the number of sparse dimensions unchanged, and NOT shrinking the size of
+  // any of the sparse dimensions.
+  // 3. When the sparse tensor has zero nnz, in which case we are free to change the shapes of
+  // both its sparse and dense dimensions.
+  //
+  // This function DOESN'T support (and will throw an error) the following cases:
+  // 1. When we attempt to change the number of sparse dimensions on a non-empty sparse tensor
+  // (such an operation will invalidate the indices stored).
+  // 2. When we attempt to change the number of dense dimensions on a non-empty sparse tensor
+  // (such an operation will behave differently from an equivalent dense tensor's resize method,
+  // and for API consistency we don't support it).
+  // 3. When we attempt to shrink the size of any of the dense dimensions on a non-empty sparse tensor
+  // (such an operation will behave differently from an equivalent dense tensor's resize method,
+  // and for API consistency we don't support it).
+  // 4. When we attempt to shrink the size of any of the sparse dimensions on a non-empty sparse tensor
+  // (this could make some of the stored indices out-of-bound and thus unsafe).
+  void resize_(int64_t sparse_dim, int64_t dense_dim, IntList size) {
+    AT_CHECK(sparse_dim + dense_dim == size.size(), "number of dimensions must be sparse_dim (", sparse_dim, ") + dense_dim (", dense_dim, "), but got ", size.size());
+    if (nnz() > 0) {
+      auto alt_options_msg = "You could try the following options:\n\
+1. If you need an empty sparse tensor of this size, call `x = torch.sparse_coo_tensor(size)`.\n\
+2. If you need to resize this tensor, you have the following options:\n\
+    1. For both sparse and dense dimensions, keep the number of them constant and the size of them non-shrinking, and then try the same call again.\n\
+    2. Or, create a new sparse tensor with the correct indices and values from this sparse tensor.";
+
+      AT_CHECK(sparse_dim == sparse_dim_,
+        "changing the number of sparse dimensions (from ", sparse_dim_, " to ", sparse_dim, ") on a non-empty sparse tensor is not supported.\n", alt_options_msg);
+
+      AT_CHECK(dense_dim == dense_dim_,
+        "changing the number of dense dimensions (from ", dense_dim_, " to ", dense_dim, ") on a non-empty sparse tensor is not supported.\n", alt_options_msg);
+
+      bool shrinking_sparse_dims = false;
+      bool shrinking_dense_dim = false;
+      auto sparse_size_original = sizes().slice(0, sparse_dim);
+      auto sparse_size_new = size.slice(0, sparse_dim);
+      for (int i = 0; i < sparse_dim; i++) {
+        if (sparse_size_new[i] < sparse_size_original[i]) {
+          shrinking_sparse_dims = true;
+          break;
+        }
+      }
+      auto dense_size_original = sizes().slice(sparse_dim);
+      auto dense_size_new = size.slice(sparse_dim);
+      for (int i = 0; i < dense_dim; i++) {
+        if (dense_size_new[i] < dense_size_original[i]) {
+          shrinking_dense_dim = true;
+          break;
+        }
+      }
+
+      AT_CHECK(!shrinking_sparse_dims,
+        "shrinking the size of sparse dimensions (from ", sparse_size_original, " to ", sparse_size_new, ") on a non-empty sparse tensor is not supported.\n", alt_options_msg);
+
+      AT_CHECK(!shrinking_dense_dim,
+        "shrinking the size of dense dimensions (from ", dense_size_original, " to ", dense_size_new, ") on a non-empty sparse tensor is not supported.\n", alt_options_msg);
+    }
+
+    if ((!size.equals(size_)) || (sparse_dim != sparse_dim_) || (dense_dim != dense_dim_)) {
+      auto nnz = values().size(0);
+      std::vector<int64_t> values_size = {nnz};
+      auto dense_size = size.slice(sparse_dim);
+      values_size.insert(values_size.end(), dense_size.begin(), dense_size.end());
+      values_.resize_(values_size);
+      indices_.resize_({sparse_dim, nnz});
+    }
+
+    size_ = size.vec();
+    sparse_dim_ = sparse_dim;
+    dense_dim_ = dense_dim;
+    refresh_numel();
   }
-  void set_values(const Tensor& values) {
-    AT_ASSERT(values.type().toSparse() == type());
-    values_ = values;
+
+  // NOTE: this function will resize the sparse tensor and also set `indices` and `values` to empty.
+  void resize_and_clear_(int64_t sparse_dim, int64_t dense_dim, IntList size) {
+    AT_CHECK(sparse_dim + dense_dim == size.size(), "number of dimensions must be sparse_dim (", sparse_dim, ") + dense_dim (", dense_dim, "), but got ", size.size());
+
+    size_ = size.vec();
+    sparse_dim_ = sparse_dim;
+    dense_dim_ = dense_dim;
+
+    auto empty_indices = at::empty({sparse_dim, 0}, indices().options());
+    std::vector<int64_t> values_size = {0};
+    auto dense_size = sizes().slice(sparse_dim);
+    values_size.insert(values_size.end(), dense_size.begin(), dense_size.end());
+    auto empty_values = at::empty(values_size, values().options());
+    set_indices_and_values_unsafe(empty_indices, empty_values);
+    refresh_numel();
   }
 
   void set_coalesced(bool coalesced) { coalesced_ = coalesced; }
-  void set_nnz(int64_t nnz) { nnz_ = nnz; }
 
+  // NOTE: this function is only used internally and not exposed to Python frontend
+  void set_nnz_and_narrow(int64_t new_nnz) {
+    AT_ASSERT(new_nnz <= nnz());
+    indices_ = indices_.narrow(1, 0, new_nnz);
+    values_ = values_.narrow(0, 0, new_nnz);
+  }
+
+  // Takes indices and values and directly puts them into the sparse tensor, no copy.
+  // NOTE: this function is unsafe because it doesn't check whether any indices are
+  // out of boundaries of `sizes`, so it should ONLY be used where we know that the
+  // indices are guaranteed to be within bounds.
   // This used to be called THSTensor_(_move)
   // NB: This used to be able to avoid a refcount bump, but I was too lazy to
   // make it happen
-  void set_indices_and_values(const Tensor& indices, const Tensor& values);
+  void set_indices_and_values_unsafe(const Tensor& indices, const Tensor& values);
+
+ private:
+  int64_t get_device_slow() const override {
+    return values_.get_device();
+  }
+
 };
 
 } // namespace at

aten/src/ATen/SparseTensorUtils.h

@@ -0,0 +1,125 @@
+#include <ATen/ATen.h>
+#include <ATen/SparseTensorImpl.h>
+
+namespace at { namespace sparse {
+
+// Just for documentary purposes
+using SparseTensor = Tensor;
+using LongTensor = Tensor;
+using IntTensor = Tensor;
+using SparseType = Type;
+
+// This is an internal utility function for getting at the SparseTensorImpl,
+// so that we can write sparse tensor specific accessors for special fields
+// in SparseTensor.  You should only use this for writing low level
+// setters/getters for SparseTensorImpl fields; otherwise, you should use
+// the low level setters/getters that were implemented using this.
+//
+// This may be called repeatedly, so make sure it's pretty cheap.
+inline SparseTensorImpl* get_sparse_impl(const SparseTensor& self) {
+  AT_ASSERTM(!self.is_variable(), "_internal_get_SparseTensorImpl: should not be a variable");
+  AT_ASSERTM(self.is_sparse(), "_internal_get_SparseTensorImpl: not a sparse tensor");
+  return static_cast<SparseTensorImpl*>(self.unsafeGetTensorImpl());
+}
+
+// Takes indices and values and directly puts them into the sparse tensor, no
+// copy.  This used to be called THSTensor_(_move)
+inline void alias_into_sparse(const SparseTensor& self, const LongTensor& indices, const Tensor& values) {
+  get_sparse_impl(self)->set_indices_and_values_unsafe(indices, values);
+}
+
+// Take indices and values and makes a (data) copy of them to put into the sparse
+// indices/values.  This used to be called THSTensor_(_set)
+inline void copy_into_sparse(const SparseTensor& self, const LongTensor& indices, const Tensor& values, bool non_blocking) {
+  alias_into_sparse(self, self._indices().type().copy(indices, non_blocking), self._values().type().copy(values, non_blocking));
+}
+
+// TODO: put this into the public API
+inline bool is_same_tensor(const Tensor& lhs, const Tensor& rhs) {
+  return lhs.unsafeGetTensorImpl() == rhs.unsafeGetTensorImpl();
+}
+
+inline bool is_same_density(const SparseTensor& self, const SparseTensor& src) {
+  return self.sparse_dim() == src.sparse_dim() && self.dense_dim() == src.dense_dim();
+}
+
+// Give us a new values tensor, with the same dimensionality
+// as 'values' but with a new number of non-zero elements.
+// TODO: Expose this for real in ATen, some day?
+// NB: Doesn't preserve data.
+inline Tensor new_values_with_size_of(const Tensor& values, int64_t nnz) {
+  std::vector<int64_t> size = values.sizes().vec();
+  size[0] = nnz;
+  return at::empty(size, values.options());
+}
+
+// NOTE [ Flatten Sparse Indices ]
+// This helper function flattens a sparse indices tensor (a LongTensor) into a 1D
+// indices tensor. E.g.,
+//   input = [[2, 4, 0],
+//            [3, 1, 10]]
+//   full_size = [2, 12]
+//   output = [ 2 * 12 + 3, 4 * 12 + 1, 0 * 12 + 10 ] = [27, 49, 10]
+//
+// In other words, assuming that each `indices[i, :]` is a valid index to a
+// tensor `t` of shape `full_size`. This returns the corresponding indices to
+// the flattened tensor `t.reshape( prod(full_size[:indices.size(0)]), -1 )`.
+// if forceClone is true, the result will forced to be a clone of self.
+// if force_clone is true, the result will forced to be a clone of self.
+inline LongTensor flatten_indices(const Tensor& indices, IntList full_size, bool force_clone = false) {
+  int64_t sparse_dim = indices.size(0);
+  if (sparse_dim == 1) {
+    if (force_clone) {
+      return indices.squeeze(0).clone();
+    } else {
+      return indices.squeeze(0);
+    }
+  } else {
+    std::vector<int64_t> indices_mult_cpu_vec;
+    indices_mult_cpu_vec.reserve(sparse_dim);
+    int64_t mult = 1;
+    for (int64_t i = sparse_dim - 1; i >= 0; i--) {
+      indices_mult_cpu_vec[i] = mult;
+      mult *= full_size[i];
+    }
+    auto indices_mult_cpu = indices.type().cpu()
+                                   .tensorFromBlob(indices_mult_cpu_vec.data(), /*size=*/{sparse_dim, 1});
+    // NB: must be blocking because this blob may be freed after this closure,
+    //     and non_blocking copy will see garbage.
+    auto indices_mult = indices_mult_cpu.to(indices.device(), /*non_blocking=*/false);
+    // Ideally we want matmul but matmul is slow on CPU Long and not implemented
+    // on CUDA Long. So mul is faster.
+    return indices.mul(indices_mult).sum(0);
+  }
+}
+
+// Flatten sparse tensor's indices from nD to 1D, similar to NOTE [ Flatten Sparse Indices ],
+// except this one allows partial flatten: only flatten on specified dims. Note that
+// the flatten indices might be uncoalesced if dims_to_flatten.size() < sparse_dim.
+// Also if input indices is already coalesced, the flattened indices will also be sorted.
+//
+// args:
+//    indices: sparse tensor indices
+//    sizes: sparse tensor sizes
+//    dims_to_flatten: a list of dim index to flatten
+//
+// Ex1:
+//   indices = [[2, 4, 0],
+//             [3, 1, 3]]
+//   sizes = [2, 12]
+//   dims_to_flatten = [0, 1]
+//   new_indices = [ 2 * 12 + 3, 4 * 12 + 1, 0 * 12 + 3 ] = [27, 49, 3]
+//
+// Ex2:
+//   dims_to_flatten = [1]
+//   new_indices = [ 3, 1, 3 ]  # uncoalesced
+inline LongTensor flatten_indices_by_dims(const LongTensor& indices, const IntList& sizes, const IntList& dims_to_flatten){
+  LongTensor new_indices = at::zeros({indices.size(1)}, indices.options());
+  for (auto d : dims_to_flatten) {
+    new_indices.mul_(sizes[d]);
+    new_indices.add_(indices.select(0, d));
+  }
+  return new_indices;
+}
+
+}} // namespace at::sparse

aten/src/ATen/Storage.h

@@ -1,41 +1,2 @@
 #pragma once
-
-#include "ATen/Scalar.h"
-
-namespace at {
-
-struct Type;
-
-struct Storage {
-  static const char RESIZABLE = 2;
-
-  Storage() {}
-  Storage(const Storage& other) = delete;
-  void operator=(const Storage&) = delete;
-
-  virtual ~Storage() {};
-  virtual size_t elementSize() const = 0;
-  virtual size_t size() const = 0;
-  virtual void* data() = 0;
-  virtual const void* data() const = 0;
-  virtual Storage& retain() = 0;
-  virtual Storage& free() = 0;
-  virtual void * unsafeGetTH(bool retain) const = 0;
-
-  virtual Storage& resize(int64_t new_size) = 0;
-
-  virtual Type & type() const = 0;
-  virtual int getDevice() const = 0;
-  virtual const char * toString() const = 0;
-
-  virtual Storage& fill(Scalar value) = 0;
-  virtual Storage& set(size_t ind, Scalar value) = 0;
-  virtual Storage& fast_set(size_t ind, Scalar value) = 0;
-  virtual Scalar get(size_t ind) = 0;
-  virtual Scalar fast_get(size_t ind) = 0;
-
-  virtual void set_flag(char flag) = 0;
-  virtual void clear_flag(char flag) = 0;
-};
-
-} // namespace at
+#include <c10/core/Storage.h>

aten/src/ATen/THLongStorageView.h

@@ -1,77 +0,0 @@
-#pragma once
-
-#include "TH/TH.h"
-#include "TH/THStorage.hpp"
-#include "TH/THTypeConversion.hpp"
-
-namespace at {
-
-enum class THLongStorageViewKind {
-  SIZE,
-  STRIDE,
-  LENGTH,
-};
-
-// make a fake storage out of a size, pointer pair...
-// used as an argument where THSize and THStride are passed into TH
-class THLongStorageView {
-public:
-  operator THLongStorage*() {
-    if (storage.size == 0 && zero_dim_to_null) {
-      return nullptr;
-    }
-    return &storage;
-  }
-
-  /*
-  // This is done as an enum, and not as static constructors, as there
-  // is no move/copy constructor for THLongStorageView
-
-  static THLongStorageView makeFromSize(ArrayRef<int64_t> ref) {
-    ...
-  }
-
-  static THLongStorageView makeFromLength(ArrayRef<int64_t> ref) {
-    ...
-  }
-  */
-
-  THLongStorageView(ArrayRef<int64_t> ref, THLongStorageViewKind kind)
-  : zero_dim_to_null(false)
-  {
-    // zero_dim_to_one converts an empty ArrayRef into [1]
-    // zero_dim_to_null converts an empty ArrayRef into a null THLongStorage
-    bool zero_dim_to_one = false;
-    bool noelem_to_empty = false;
-    switch (kind) {
-      case THLongStorageViewKind::SIZE:
-        zero_dim_to_one = true;
-        break;
-      case THLongStorageViewKind::STRIDE:
-        zero_dim_to_null = true;
-        break;
-      case THLongStorageViewKind::LENGTH:
-        break;
-    }
-
-    if(zero_dim_to_one && ref.size() == 0) {
-      // make storage of size 0 actually a 1-length storage with 1 element
-      // so that our 0-dim tensors get allocated as 1-dim inside TH
-      one = 1;
-      storage.data_ptr = {&one, kCPU}; // non-owning
-      storage.size = 1;
-    } else {
-      storage.data_ptr = {const_cast<void*>(static_cast<const void*>(ref.data())), kCPU}; // non-owning
-      storage.size = ref.size();
-    }
-    storage.scalar_type = at::CTypeToScalarType<th::from_type<int64_t>>::to();
-    storage.refcount = 0;
-    storage.flag = 0;
-  }
-private:
-  int64_t one;
-  THLongStorage storage;
-  bool zero_dim_to_null;
-};
-
-}

aten/src/ATen/Tensor.h

@@ -0,0 +1,2 @@
+#pragma once
+#include <ATen/core/Tensor.h>

aten/src/ATen/TensorAccessor.h

@@ -1,51 +1,2 @@
 #pragma once
-
-#include <cstddef>
-#include <stdint.h>
-
-#include "ATen/ScalarType.h"
-
-namespace at {
-
-
-template<typename T, size_t N>
-class TensorAccessorBase {
-public:
-  TensorAccessorBase(T * data_, const int64_t * sizes_, const int64_t * strides_)
-  : data_(data_), sizes_(sizes_), strides_(strides_) {}
-  IntList sizes() {
-    return IntList(sizes_,N);
-  }
-  IntList strides() {
-    return IntList(strides_,N);
-  }
-  int64_t stride(int64_t i) { return strides()[i]; }
-  int64_t size(int64_t i) { return sizes()[i]; }
-protected:
-  T * data_;
-  const int64_t* sizes_;
-  const int64_t* strides_;
-};
-
-template<typename T, size_t N>
-class TensorAccessor : public TensorAccessorBase<T,N> {
-public:
-  TensorAccessor(T * data_, const int64_t * sizes_, const int64_t * strides_)
-  : TensorAccessorBase<T,N>(data_,sizes_,strides_) {}
-
-  TensorAccessor<T,N-1> operator[](int64_t i) {
-    return TensorAccessor<T,N-1>(this->data_ + this->strides_[0]*i,this->sizes_+1,this->strides_+1);
-  }
-};
-
-template<typename T>
-class TensorAccessor<T,1> : public TensorAccessorBase<T,1> {
-public:
-  TensorAccessor(T * data_, const int64_t * sizes_, const   int64_t * strides_)
-  : TensorAccessorBase<T,1>(data_,sizes_,strides_) {}
-  T & operator[](int64_t i) {
-    return this->data_[this->strides_[0]*i];
-  }
-};
-
-}
+#include <ATen/core/TensorAccessor.h>

aten/src/ATen/TensorBase.h

@@ -1,108 +0,0 @@
-#pragma once
-
-#include "ATen/TensorImpl.h"
-#include "ATen/UndefinedTensor.h"
-
-namespace at { namespace detail {
-
-// TensorBaseImpl is the base class for Tensor which handles the reference counting
-template<bool is_strong>
-struct TensorBaseImpl {
-  TensorBaseImpl(): TensorBaseImpl(UndefinedTensor::singleton(), false) {}
-  TensorBaseImpl(TensorImpl * self, bool should_retain)
-  : pImpl(self) {
-    if (pImpl == nullptr) {
-      throw std::runtime_error("TensorBaseImpl with nullptr not supported");
-    }
-    if(should_retain && pImpl != UndefinedTensor::singleton()) {
-      retain();
-    }
-  }
-  TensorBaseImpl(const TensorBaseImpl & rhs)
-  : pImpl(rhs.pImpl) {
-    if (pImpl != UndefinedTensor::singleton()) {
-      retain();
-    }
-  }
-  TensorBaseImpl(TensorBaseImpl && rhs) noexcept
-  : pImpl(rhs.pImpl) {
-    rhs.pImpl = UndefinedTensor::singleton();
-  }
-  ~TensorBaseImpl() {
-    if (pImpl != UndefinedTensor::singleton()) {
-      release();
-    }
-  }
-  TensorBaseImpl & operator=(TensorBaseImpl && rhs) & {
-    rhs.swap(*this);
-    return *this;
-  }
-  TensorBaseImpl & operator=(TensorBaseImpl const & rhs) & {
-    //TensorBaseImpl ctor retains original rhs.pImpl
-    //then rhs.pImpl is swapped with this->pImpl
-    //finally TensorBaseImpl dtor releases rhs.pImpl, which was originally this->pImpl
-    TensorBaseImpl(rhs).swap(*this);
-    return *this;
-  }
-  int64_t dim() const {
-    if (is_strong) {
-      return pImpl->dim();
-    } else {
-      AT_ERROR("Can't call dim() on a WeakTensor");
-    }
-  }
-  void reset() {
-    TensorBaseImpl().swap(*this);
-  }
-  void reset(TensorImpl * rhs) {
-    TensorBaseImpl(rhs, true).swap(*this);
-  }
-  void reset(TensorImpl * rhs, bool should_retain) {
-    TensorBaseImpl(rhs, should_retain).swap(*this );
-  }
-  void swap(TensorBaseImpl & rhs) {
-    TensorImpl * tmp = pImpl;
-    pImpl = rhs.pImpl;
-    rhs.pImpl = tmp;
-  }
-  TensorImpl * get() const {
-    return pImpl;
-  }
-  TensorImpl * detach() {
-    TensorImpl * ret = pImpl;
-    pImpl = UndefinedTensor::singleton();
-    return ret;
-  }
-
-  bool defined() const {
-    return pImpl != UndefinedTensor::singleton();
-  }
-
-  friend struct Type;
-
-  //TODO(zach): sort out friend structes
-public:
-  TensorImpl * pImpl;
-
-private:
-  void retain() {
-    if (is_strong) {
-      pImpl->retain();
-    } else {
-      pImpl->weak_retain();
-    }
-  }
-
-  void release() {
-    if (is_strong) {
-      pImpl->release();
-    } else {
-      pImpl->weak_release();
-    }
-  }
-};
-
-using TensorBase = TensorBaseImpl<true>;
-using WeakTensorBase = TensorBaseImpl<false>;
-
-}} // namespace at::detail

aten/src/ATen/TensorGeometry.cpp

@@ -1,23 +1,15 @@
 #include <ATen/TensorGeometry.h>
+#include <ATen/TensorUtils.h>
 
 #include <ATen/ATen.h>
 
 namespace at {
 
 bool TensorGeometry::is_contiguous() const {
-  int64_t dim = sizes_.size();
-  int64_t expected_stride = 1;
-  for (int64_t i = dim - 1; i >= 0; i--) {
-    if (sizes_[i] != 1 && strides_[i] != expected_stride) {
-      return false;
-    }
-    expected_stride *= sizes_[i];
+  if (numel_ == 0) {
+    return true;
   }
-  return true;
-}
-
-Tensor TensorGeometry::zeros_with_stride(const Type& type) const {
-  return type.tensor(sizes_, strides_).zero_();
+  return at::geometry_is_contiguous(sizes_, strides_);
 }
 
 } // namespace at

aten/src/ATen/TensorGeometry.h

@@ -5,11 +5,11 @@
 
 namespace at {
 
-struct AT_API TensorGeometry {
+struct CAFFE2_API TensorGeometry {
   TensorGeometry() : storage_offset_(0) {}
 
   explicit TensorGeometry(IntList sizes)
-    : sizes_(sizes)
+    : sizes_(sizes.vec())
     , strides_(sizes.size())
     , storage_offset_(0) {
       int64_t dim = sizes.size();
@@ -18,19 +18,18 @@ struct AT_API TensorGeometry {
         strides_[i] = expected_stride;
         expected_stride *= sizes_[i];
       }
+      numel_ = expected_stride;
   }
 
   explicit TensorGeometry(const Tensor& t)
-    : sizes_(t.sizes())
-    , strides_(t.strides())
-    , storage_offset_(t.storage_offset()) {}
+    : sizes_(t.sizes().vec())
+    , strides_(t.strides().vec())
+    , storage_offset_(t.storage_offset())
+    , numel_(t.numel()) {}
 
   // true if the tensor is contiguous
   bool is_contiguous() const;
 
-  // creates a new tensor with the sizes and strides of the source
-  Tensor zeros_with_stride(const Type& type) const;
-
   int64_t dim() const { return sizes_.size(); }
   int64_t size(int64_t dim) const {
     dim = maybe_wrap_dim(dim, this->dim());
@@ -43,13 +42,7 @@ struct AT_API TensorGeometry {
   }
   IntList strides() const { return IntList{ strides_ }; }
   int64_t storage_offset() const { return storage_offset_; }
-  int64_t numel() const {
-    int64_t r = 1;
-    for (auto s : sizes()) {
-      r *= s;
-    }
-    return r;
-  }
+  int64_t numel() const { return numel_; }
 
   TensorGeometry transpose(int64_t dim0, int64_t dim1) {
     TensorGeometry r = *this; // copy
@@ -63,6 +56,7 @@ struct AT_API TensorGeometry {
   std::vector<int64_t> sizes_;
   std::vector<int64_t> strides_;
   int64_t storage_offset_;
+  int64_t numel_;
 };
 
 } // namespace at

aten/src/ATen/TensorImpl.cpp

@@ -1,36 +0,0 @@
-#include <ATen/TensorImpl.h>
-
-#include <ATen/Tensor.h>
-#include <ATen/optional.h>
-
-namespace at {
-Tensor& TensorImpl::grad() {
-  AT_ERROR("grad is not implemented for Tensor");
-}
-
-const Tensor& TensorImpl::grad() const {
-  AT_ERROR("grad is not implemented for Tensor");
-}
-
-Tensor TensorImpl::detach() const {
-  AT_ERROR("detach is not implemented for Tensor");
-}
-
-void TensorImpl::backward(
-    at::optional<Tensor> gradient,
-    bool keep_graph,
-    bool create_graph) {
-  AT_ERROR("backward is not implemented for Tensor");
-}
-
-void TensorImpl::set_data(Tensor new_data) {
-  AT_ERROR("set_type is not implemented for Tensor");
-}
-
-void Tensor::backward(
-    at::optional<Tensor> gradient,
-    bool keep_graph,
-    bool create_graph) {
-  pImpl->backward(std::move(gradient), keep_graph, create_graph);
-}
-} // namespace at

aten/src/ATen/TensorImpl.h

@@ -1,98 +0,0 @@
-#pragma once
-
-#include <atomic>
-#include <memory>
-
-#include "ATen/Retainable.h"
-#include "ATen/ScalarType.h"
-#include "ATen/optional.h"
-
-namespace at {
-class Scalar;
-struct Type;
-struct Storage;
-struct Tensor;
-} // namespace at
-
-namespace at {
-struct TensorImpl : public Retainable {
-  explicit TensorImpl(Type * type)
-  : is_scalar(false), type_(type) {}
-
-  Type & type() const {
-    return *type_;
-  }
-  virtual const char * toString() const = 0;
-  virtual IntList sizes() const = 0;
-  virtual IntList strides() const = 0;
-  virtual int64_t dim() const = 0;
-  /**
-   * Perform a conversion of this tensor to a scalar, if numel() == 1.
-   * Otherwise, raise an error.
-   */
-  virtual Scalar localScalar() = 0;
-  virtual void * unsafeGetTH(bool retain) = 0;
-  virtual std::unique_ptr<Storage> storage() = 0;
-  friend struct Type;
-
-  int64_t numel() {
-    int64_t n = 1;
-    for (auto s : sizes()) {
-      n *= s;
-    }
-    return n;
-  }
-
-  // 0-dim patchup of TH requires us to have a flag marking
-  // if a Tensor should be treated as 0-dim.
-  // the generated wrapper manipulates this flag.
-  // the setter should never be exposed in Tensor's public API
-  // because eventually we would like isScalar() to just be dim() == 0;
-  bool isScalar() const {
-    return is_scalar;
-  }
-  // this is called by the generated wrapper code when there are conditions
-  // when this output tensor should be a scalar. e.g. when all inputs
-  // to a function 'add' were scalars, then condition_when_scalar == true.
-  // we also prevent this from getting marked as a scalar if it is not
-  // the right shape afterall.
-  TensorImpl* maybeScalar(bool condition_when_scalar) {
-    is_scalar = false; //force dim() to tell the truth for TH
-    is_scalar = condition_when_scalar && dim() == 1 && sizes()[0] == 1;
-    return this;
-  }
-  void setScalar(bool s) {
-    is_scalar = s;
-  }
-
-  // ~~~~~ Autograd API ~~~~~
-  // Some methods below are defined in TensorImpl.cpp because Tensor is an
-  // incomplete type.
-
-  AT_API virtual void set_requires_grad(bool requires_grad) {
-    AT_ERROR("set_requires_grad is not implemented for Tensor");
-  }
-  AT_API virtual bool requires_grad() const {
-    AT_ERROR("requires_grad is not implemented for Tensor");
-  }
-
-  AT_API virtual Tensor& grad();
-  AT_API virtual const Tensor& grad() const;
-
-  AT_API virtual Tensor detach() const;
-  AT_API virtual void detach_() {
-    AT_ERROR("detach_ is not implemented for Tensor");
-  }
-
-  AT_API virtual void backward(
-      at::optional<Tensor> gradient,
-      bool keep_graph,
-      bool create_graph);
-
-  AT_API virtual void set_data(Tensor new_data);
-
-protected:
-  bool is_scalar;
-  Type * type_;
-};
-} // namespace at

aten/src/ATen/TensorOperators.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include "ATen/Scalar.h"
+#include <c10/core/Scalar.h>
 #include "ATen/Tensor.h"
 #include "ATen/Type.h"
 
@@ -9,7 +9,12 @@
 
 namespace at {
 
-
+inline Tensor & Tensor::operator=(Tensor const & rhs) && {
+  return copy_(rhs);
+}
+inline Tensor & Tensor::operator=(Tensor && rhs) && {
+  return copy_(rhs);
+}
 inline Tensor & Tensor::operator=(Scalar v) && {
   return fill_(v);
 }
@@ -42,9 +47,8 @@ inline Tensor& Tensor::operator/=(Scalar other) {
 }
 inline Tensor Tensor::operator[](Scalar index) const {
   AT_CHECK(
-      index.local().isIntegral(),
-      "Can only index tensors with integral scalars (got ",
-      index.toTensor().type().toString(), ")");
+      index.isIntegral(),
+      "Can only index tensors with integral scalars");
   return select(0, index.toLong());
 }
 inline Tensor Tensor::operator[](Tensor index) const {
@@ -55,7 +59,7 @@ inline Tensor Tensor::operator[](Tensor index) const {
       index.dim() == 0,
       "Can only index with tensors that are scalars (zero-dim)");
   // The Scalar(Tensor) constructor is explicit, so we need to call it.
-  return this->operator[](Scalar(index));
+  return this->operator[](index.item());
 }
 inline Tensor Tensor::operator[](int64_t index) const {
   return select(0, index);
@@ -64,9 +68,9 @@ inline Tensor Tensor::operator[](int64_t index) const {
 #define AT_FORALL_BINARY_OPS(_) \
 _(+,x.add(y), y.add(x)) \
 _(*,x.mul(y), y.mul(x)) \
-_(-,x.sub(y), y.type().tensor().resize_(y.sizes()).fill_(x).sub_(y)) \
-_(/,x.div(y), y.type().tensor().resize_(y.sizes()).fill_(x).div_(y)) \
-_(%,x.remainder(y), y.type().tensor().resize_(y.sizes()).fill_(x).remainder_(y)) \
+_(-,x.sub(y), ::at::empty(y.sizes(), y.options()).fill_(x).sub_(y)) \
+_(/,x.div(y), ::at::empty(y.sizes(), y.options()).fill_(x).div_(y)) \
+_(%,x.remainder(y), ::at::empty(y.sizes(), y.options()).fill_(x).remainder_(y)) \
 _(<,x.lt(y), y.gt(x)) \
 _(<=,x.le(y), y.ge(x)) \
 _(>,x.gt(y),y.lt(x)) \

aten/src/ATen/TensorOptions.cpp

@@ -1,19 +0,0 @@
-#include <ATen/TensorOptions.h>
-
-#include <ATen/Device.h>
-#include <ATen/Layout.h>
-#include <ATen/OptionsGuard.h>
-#include <ATen/ScalarType.h>
-#include <ATen/optional.h>
-
-namespace at {
-
-TensorOptions::TensorOptions(bool use_thread_local_default_options) {
-  if (use_thread_local_default_options) {
-    this->dtype(DefaultTensorOptions::get().dtype());
-    this->device(DefaultTensorOptions::get().device());
-    this->layout(DefaultTensorOptions::get().layout());
-    this->requires_grad(DefaultTensorOptions::get().requires_grad());
-  }
-}
-} // namespace at

aten/src/ATen/TensorOptions.h

@@ -1,279 +1,2 @@
 #pragma once
-
-#include <ATen/Context.h>
-#include <ATen/Device.h>
-#include <ATen/DeviceGuard.h>
-#include <ATen/Layout.h>
-#include <ATen/ScalarType.h>
-#include <ATen/Tensor.h>
-#include <ATen/Type.h>
-
-#include <cstddef>
-#include <utility>
-
-namespace at {
-
-/// A class to encapsulate construction axes of a `Tensor`.
-/// `TensorOptions` is a virtual class to enable overriding of certain methods
-/// by subclasses in other libraries, such as PyTorch. In PyTorch, there is a
-/// `torch::TensorOptions` subclass of this `TensorOptions`, which changes
-/// `type()` to return a variable type instead of a tensor type, such that
-/// variables are created inside factory methods, instead of tensors.
-struct TensorOptions {
-  TensorOptions() : TensorOptions(/*use_thread_local_default_options=*/true) {}
-
-  /// Constructs the `TensorOptions` with defaults taken from the thread local
-  /// `TensorOptions` object if `use_thread_local_default_options`, else
-  /// defaults to:
-  /// - dtype: kFloat,
-  /// - device: kCPU,
-  /// - layout: kStrided,
-  /// - requires_grad: false
-  explicit TensorOptions(bool use_thread_local_default_options);
-
-  /// Constructs the `TensorOptions` from the type of the given `Tensor`.
-  /// If the `Tensor` has a CUDA type, the `device_index` will match that of the
-  /// tensor. The `requires_grad` property of the tensor is ignored and set to
-  /// false in the created `TensorOptions`.  See the constructor from `Type` for
-  /// the semantics w.r.t. the `type()` method.
-  explicit TensorOptions(Tensor tensor, bool discard_runtime_type = false) {
-    if (!discard_runtime_type) {
-      type_ = &tensor.type();
-    }
-    this->dtype(tensor.dtype());
-    this->device(tensor.device());
-    this->layout(tensor.layout());
-  }
-
-  /// Constructs the `TensorOptions` from a type and a `device_index`.
-  ///
-  /// If `discard_runtime_type` is false (the default), the behavior of
-  /// `TensorOptions::type()` is changed in that it will always return this
-  /// `type`, irrespective of any `device` or `dtype` or `layout` specified at a
-  /// later time. This is to ensure that when a `TensorOptions` object is
-  /// constructed from a tensor's type, and that type has a dynamic type other
-  /// than `at::Type` (e.g. `torch::autograd::VariableType`), constructing a new
-  /// tensor from this `TensorOptions` will use this same derived type. If
-  /// instead the given `type` were destructured into its components (backend,
-  /// dtype and layout), information about the runtime type of the `Type` would
-  /// be lost. Set `discard_runtime_type` to `true` to always destructure the
-  /// type into its components and discard its runtime type.
-  /* implicit */ TensorOptions(
-      const Type& type,
-      int32_t device_index = -1,
-      bool discard_runtime_type = false) {
-    if (!discard_runtime_type) {
-      type_ = &type;
-    }
-    this->dtype(type.scalarType());
-    this->device({type.backend(), device_index});
-    this->layout(type.layout());
-  }
-
-  /// Constructs a `TensorOptions` object with the given layout.
-  /* implicit */ TensorOptions(Layout layout) : TensorOptions() {
-    this->layout(layout);
-  }
-
-  /// Constructs a `TensorOptions` object with the given device.
-  /* implicit */ TensorOptions(Device device) : TensorOptions() {
-    this->device(device);
-  }
-
-  /// Constructs a `TensorOptions` object from a backend, forwarded to the
-  /// `Device` constructor.
-  /* implicit */ TensorOptions(Backend backend)
-      : TensorOptions(Device(backend)) {}
-
-  /// Constructs a `TensorOptions` object with the given dtype.
-  /* implicit */ TensorOptions(ScalarType dtype) : TensorOptions() {
-    this->dtype(dtype);
-  }
-
-  /// True if all elements of the `TensorOptions` match that of the other.
-  bool operator==(const TensorOptions& other) const noexcept {
-    return dtype_ == other.dtype_ && layout_ == other.layout_ &&
-        device_ == other.device_ && requires_grad_ == other.requires_grad_;
-  }
-
-  /// True if any of the elements of this `TensorOptions` do not match that of
-  /// the other.
-  bool operator!=(const TensorOptions& other) const noexcept {
-    return !(*this == other);
-  }
-
-  /// Discards the runtime type stored if the `TensorOptions` was constructed
-  /// from a `Tensor` or a `Type`. See the documentation of the constructor from
-  /// a `Type` for implications on the behavior of the `type()` method on
-  /// `TensorOptions`.
-  const TensorOptions& discard_runtime_type() const {
-    type_ = nullptr;
-    return *this;
-  }
-
-  /// Sets the device of the `TensorOptions`.
-  TensorOptions& device(Device device) {
-    device_ = std::move(device);
-    update_underlying_type();
-    return *this;
-  }
-
-  /// Sets the device of the `TensorOptions` to CUDA, and then sets the device
-  /// index to the given one.
-  TensorOptions& device_index(int32_t device_index) {
-    return device({Device::Type::CUDA, device_index});
-  }
-
-  /// Sets the dtype of the `TensorOptions`.
-  TensorOptions& dtype(ScalarType dtype) {
-    dtype_ = dtype;
-    update_underlying_type();
-    return *this;
-  }
-
-  /// Sets the layout of the `TensorOptions`.
-  TensorOptions& layout(Layout layout) {
-    layout_ = layout;
-    update_underlying_type();
-    return *this;
-  }
-
-  /// Sets the `requires_grad` property of the `TensorOptions`.
-  TensorOptions& requires_grad(bool requires_grad) {
-    requires_grad_ = requires_grad;
-    return *this;
-  }
-
-  /// Returns the device of the `TensorOptions`.
-  const Device& device() const noexcept {
-    return device_;
-  }
-
-  /// Returns the device index of the `TensorOptions`.
-  int32_t device_index() const noexcept {
-    return device_.index();
-  }
-
-  /// Returns the dtype of the `TensorOptions`.
-  ScalarType dtype() const noexcept {
-    return dtype_;
-  }
-
-  /// Returns the layout of the `TensorOptions`.
-  Layout layout() const noexcept {
-    return layout_;
-  }
-
-  /// Returns the `requires_grad` property of the `TensorOptions`.
-  bool requires_grad() const noexcept {
-    return requires_grad_;
-  }
-
-  /// Constructs an `at::Type` from the members of the `TensorOptions`.
-  const Type& type() const {
-    if (type_ != nullptr) {
-      return *type_;
-    }
-    return getType(backend(), dtype_);
-  }
-
- private:
-  /// Updates any stored underlying type to the current construction axes.
-  void update_underlying_type() {
-    if (type_) {
-      type_ = &type_->toScalarType(dtype_).toBackend(backend());
-    }
-  }
-
-  // Resolves the ATen backend specified by the current construction axes.
-  Backend backend() const noexcept {
-    Backend backend;
-    if (device_.type() == Device::Type::CPU) {
-      backend = (layout_ == kStrided) ? kCPU : kSparseCPU;
-    } else {
-      backend = (layout_ == kStrided) ? kCUDA : kSparseCUDA;
-    }
-    return backend;
-  }
-
- private:
-  ScalarType dtype_{kFloat};
-  Device device_{Device::Type::CPU};
-  Layout layout_{Layout::Strided};
-  bool requires_grad_{false};
-  // Not part of the observable API, so make `mutable` so we can set it to
-  // `null` in `discard_runtime_type`.
-  mutable const Type* type_{nullptr};
-};
-
-/// Convenience function that returns a `TensorOptions` object with the `dtype`
-/// set to the given one.
-inline TensorOptions dtype(ScalarType dtype) {
-  return TensorOptions().dtype(dtype);
-}
-
-/// Convenience function that returns a `TensorOptions` object with the `layout`
-/// set to the given one.
-inline TensorOptions layout(Layout layout) {
-  return TensorOptions().layout(layout);
-}
-
-/// Convenience function that returns a `TensorOptions` object with the `device`
-/// set to the given one.
-inline TensorOptions device(Device device) {
-  return TensorOptions().device(std::move(device));
-}
-
-/// Convenience function that returns a `TensorOptions` object with the
-/// `device_index` set to the given one.
-inline TensorOptions device_index(int32_t device_index) {
-  return TensorOptions().device_index(device_index);
-}
-
-/// Convenience function that returns a `TensorOptions` object with the
-/// `requires_grad` set to the given one.
-inline TensorOptions requires_grad(bool requires_grad = true) {
-  return TensorOptions().requires_grad(requires_grad);
-}
-
-/// From Tensor.h
-inline TensorOptions Tensor::options() const {
-  return TensorOptions(*this);
-}
-
-namespace detail {
-inline Tensor to(
-    const Tensor& tensor,
-    const TensorOptions& options,
-    bool non_blocking) {
-  // Don't copy if the options match.
-  if (tensor.options() == options) {
-    return tensor;
-  }
-  DeviceGuard guard(options.device());
-  return options.type().copy(tensor, non_blocking);
-}
-} // namespace detail
-
-inline Tensor Tensor::to(Device device, ScalarType dtype, bool non_blocking)
-    const {
-  if (this->device() == device && this->dtype() == dtype) {
-    return *this;
-  }
-  return detail::to(*this, options().device(device).dtype(dtype), non_blocking);
-}
-
-inline Tensor Tensor::to(ScalarType dtype, bool non_blocking) const {
-  if (this->dtype() == dtype) {
-    return *this;
-  }
-  return detail::to(*this, options().dtype(dtype), non_blocking);
-}
-
-inline Tensor Tensor::to(Device device, bool non_blocking) const {
-  if (this->device() == device) {
-    return *this;
-  }
-  return detail::to(*this, options().device(device), non_blocking);
-}
-} // namespace at
+#include <ATen/core/TensorOptions.h>