2022-09-14 nightly release (55ca297)

pytorch · ezyang · Sep 11, 2022 · Sep 12, 2022 · Sep 13, 2022 · Sep 14, 2022
commit f35657efc9d8a643c54d2232518c48dabca93eed
diff --git a/.circleci/docker/build.sh b/.circleci/docker/build.sh
@@ -379,7 +379,7 @@ docker build \
        --build-arg "NINJA_VERSION=${NINJA_VERSION:-}" \
        --build-arg "KATEX=${KATEX:-}" \
        --build-arg "ROCM_VERSION=${ROCM_VERSION:-}" \
-       --build-arg "PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH:-gfx900;gfx906}" \
+       --build-arg "PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH:-gfx906}" \
        --build-arg "IMAGE_NAME=${IMAGE_NAME}" \
        --build-arg "UCX_COMMIT=${UCX_COMMIT}" \
        --build-arg "UCC_COMMIT=${UCC_COMMIT}" \

diff --git a/.circleci/docker/common/install_cudnn.sh b/.circleci/docker/common/install_cudnn.sh
@@ -4,13 +4,7 @@ if [[ ${CUDNN_VERSION} == 8 ]]; then
     # cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
     mkdir tmp_cudnn && cd tmp_cudnn
     CUDNN_NAME="cudnn-linux-x86_64-8.3.2.44_cuda11.5-archive"
-    if [[ ${CUDA_VERSION:0:4} == "11.7" ]]; then
-        CUDNN_NAME="cudnn-linux-x86_64-8.5.0.96_cuda11-archive"
-        curl -OLs https://ossci-linux.s3.amazonaws.com/${CUDNN_NAME}.tar.xz
-    else
-        curl -OLs  https://developer.download.nvidia.com/compute/redist/cudnn/v8.3.2/local_installers/11.5/${CUDNN_NAME}.tar.xz
-    fi
-
+    curl -OLs  https://developer.download.nvidia.com/compute/redist/cudnn/v8.3.2/local_installers/11.5/${CUDNN_NAME}.tar.xz
     tar xf ${CUDNN_NAME}.tar.xz
     cp -a ${CUDNN_NAME}/include/* /usr/include/
     cp -a ${CUDNN_NAME}/include/* /usr/local/cuda/include/

diff --git a/.circleci/docker/ubuntu-cuda/Dockerfile b/.circleci/docker/ubuntu-cuda/Dockerfile
@@ -118,7 +118,6 @@ COPY --from=pytorch/llvm:9.0.1 /opt/llvm /opt/llvm
 
 # Install CUDNN
 ARG CUDNN_VERSION
-ARG CUDA_VERSION
 COPY ./common/install_cudnn.sh install_cudnn.sh
 RUN if [ "${CUDNN_VERSION}" -eq 8 ]; then bash install_cudnn.sh; fi
 RUN rm install_cudnn.sh

diff --git a/.circleci/scripts/windows_cudnn_install.sh b/.circleci/scripts/windows_cudnn_install.sh
@@ -18,7 +18,7 @@ case ${CUDA_VERSION} in
         ;;
     11.7)
         # Use cudnn8.3 with hard-coded cuda11.5 version
-        cudnn_file_name="cudnn-windows-x86_64-8.5.0.96_cuda11-archive"
+        cudnn_file_name="cudnn-windows-x86_64-8.3.2.44_cuda11.5-archive"
         ;;
     *)
         echo "CUDA_VERSION: ${CUDA_VERSION} not supported yet"

diff --git a/.github/ci_commit_pins/torchdynamo.txt b/.github/ci_commit_pins/torchdynamo.txt
@@ -1 +1 @@
-ffd056dc1510bdfecafb689ed87601055694f3e6
+138eb96c36c6ba5a14e65d28fad3cbd5a6e6dcfc
@@ -1 +1 @@
-a89b1957a62e2f68f001d5d60268743edbe164d8
+ebb68f337de8f66df5bb83f8266f3a9c7020094e
@@ -323,6 +323,7 @@
   - '*'
   approved_by:
   - pytorch/metamates
+  - mruberry
   mandatory_checks_name:
   - Facebook CLA Check
   - Lint

diff --git a/.jenkins/pytorch/build.sh b/.jenkins/pytorch/build.sh
@@ -68,8 +68,13 @@ fi
 pip_install -r requirements.txt || true
 
 # Enable LLVM dependency for TensorExpr testing
-export USE_LLVM=/opt/llvm
-export LLVM_DIR=/opt/llvm/lib/cmake/llvm
+if [[ "$BUILD_ENVIRONMENT" == *rocm* ]]; then
+  export USE_LLVM=/opt/rocm/llvm
+  export LLVM_DIR=/opt/rocm/llvm/lib/cmake/llvm
+else
+  export USE_LLVM=/opt/llvm
+  export LLVM_DIR=/opt/llvm/lib/cmake/llvm
+fi
 
 # TODO: Don't install this here
 if ! which conda; then
@@ -146,9 +151,9 @@ if [[ "$BUILD_ENVIRONMENT" == *rocm* ]]; then
   fi
 
   if [[ -n "$CI" && -z "$PYTORCH_ROCM_ARCH" ]]; then
-      # Set ROCM_ARCH to gfx900 and gfx906 for CI builds, if user doesn't override.
-      echo "Limiting PYTORCH_ROCM_ARCH to gfx90[06] for CI builds"
-      export PYTORCH_ROCM_ARCH="gfx900;gfx906"
+      # Set ROCM_ARCH to gfx906 for CI builds, if user doesn't override.
      echo "Limiting PYTORCH_ROCM_ARCH to gfx906 for CI builds"
+      export PYTORCH_ROCM_ARCH="gfx906"
   fi
 
   # hipify sources

diff --git a/.jenkins/pytorch/common_utils.sh b/.jenkins/pytorch/common_utils.sh
@@ -141,12 +141,6 @@ function checkout_install_torchdynamo() {
   popd
 }
 
-function install_functorch() {
-  pushd functorch
-  time python setup.py develop
-  popd
-}
-
 function test_functorch() {
   python test/run_test.py --functorch --verbose
 }

diff --git a/.jenkins/pytorch/macos-test.sh b/.jenkins/pytorch/macos-test.sh
@@ -177,7 +177,6 @@ test_dynamo() {
 }
 
 if [[ "${TEST_CONFIG}" == *functorch* ]]; then
-  install_functorch
   test_functorch
 elif [[ $NUM_TEST_SHARDS -gt 1 ]]; then
   test_python_shard "${SHARD_NUMBER}"

diff --git a/.jenkins/pytorch/test.sh b/.jenkins/pytorch/test.sh
@@ -180,9 +180,6 @@ test_dynamo_shard() {
     echo "NUM_TEST_SHARDS must be defined to run a Python test shard"
     exit 1
   fi
-  pushd functorch
-  python setup.py develop
-  popd
   # Temporarily disable test_fx for dynamo pending the investigation on TTS
   # regression in https://github.com/pytorch/torchdynamo/issues/784
   time python test/run_test.py \
@@ -332,6 +329,14 @@ test_libtorch() {
 
 test_aot_compilation() {
   echo "Testing Ahead of Time compilation"
+  ln -sf "$TORCH_LIB_DIR"/libc10* "$TORCH_BIN_DIR"
+  ln -sf "$TORCH_LIB_DIR"/libtorch* "$TORCH_BIN_DIR"
+
+  # Make test_reports directory
+  # NB: the ending test_libtorch must match the current function name for the current
+  # test reporting process (in print_test_stats.py) to function as expected.
+  TEST_REPORTS_DIR=test/test-reports/cpp-unittest/test_aot_compilation
+  mkdir -p $TEST_REPORTS_DIR
   if [ -f "$TORCH_BIN_DIR"/test_mobile_nnc ]; then "$TORCH_BIN_DIR"/test_mobile_nnc --gtest_output=xml:$TEST_REPORTS_DIR/test_mobile_nnc.xml; fi
   # shellcheck source=test/mobile/nnc/test_aot_compile.sh
   if [ -f "$TORCH_BIN_DIR"/aot_model_compiler_test ]; then source test/mobile/nnc/test_aot_compile.sh; fi
@@ -686,7 +691,6 @@ elif [[ "${BUILD_ENVIRONMENT}" == *-mobile-lightweight-dispatch* ]]; then
 elif [[ "${TEST_CONFIG}" = docs_test ]]; then
   test_docs_test
 elif [[ "${TEST_CONFIG}" == *functorch* ]]; then
-  install_functorch
   test_functorch
 else
   install_torchvision

diff --git a/.jenkins/pytorch/win-test-helpers/build_pytorch.bat b/.jenkins/pytorch/win-test-helpers/build_pytorch.bat
@@ -144,7 +144,7 @@ python setup.py install --cmake && sccache --show-stats && (
   if "%BUILD_ENVIRONMENT%"=="" (
     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 %TMP_DIR_WIN%\%IMAGE_COMMIT_TAG%.7z %CONDA_PARENT_DIR%\Miniconda3\Lib\site-packages\torch %CONDA_PARENT_DIR%\Miniconda3\Lib\site-packages\torchgen %CONDA_PARENT_DIR%\Miniconda3\Lib\site-packages\caffe2 && copy /Y "%TMP_DIR_WIN%\%IMAGE_COMMIT_TAG%.7z" "%PYTORCH_FINAL_PACKAGE_DIR%\"
+    7z a %TMP_DIR_WIN%\%IMAGE_COMMIT_TAG%.7z %CONDA_PARENT_DIR%\Miniconda3\Lib\site-packages\torch %CONDA_PARENT_DIR%\Miniconda3\Lib\site-packages\torchgen %CONDA_PARENT_DIR%\Miniconda3\Lib\site-packages\caffe2 %CONDA_PARENT_DIR%\Miniconda3\Lib\site-packages\functorch && copy /Y "%TMP_DIR_WIN%\%IMAGE_COMMIT_TAG%.7z" "%PYTORCH_FINAL_PACKAGE_DIR%\"
     if errorlevel 1 exit /b
     if not errorlevel 0 exit /b
 

diff --git a/.jenkins/pytorch/win-test-helpers/install_test_functorch.bat b/.jenkins/pytorch/win-test-helpers/install_test_functorch.bat
@@ -6,15 +6,6 @@ if not errorlevel 0 (
   exit /b
 )
 
-pushd functorch
-echo "Install functorch"
-:: --no-deps because for some reason, on windows, `torch` isn't found in
-:: `pip list` despite being installed. With just `python setup.py develop`,
-:: setuptools explicitly checks for the existence of torch and can't find it.
-python setup.py develop --no-deps
-popd
-if ERRORLEVEL 1 goto fail
-
 echo "Installing test dependencies"
 pip install networkx
 if errorlevel 1 exit /b

diff --git a/.lintrunner.toml b/.lintrunner.toml
@@ -318,6 +318,7 @@ exclude_patterns = [
     'aten/src/ATen/native/vulkan/api/vk_mem_alloc.h',
     'test/cpp/jit/upgrader_models/*.ptl',
     'test/cpp/jit/upgrader_models/*.ptl.ff',
+    'cmake/External/nccl.patch',
 ]
 command = [
     'python3',
@@ -347,6 +348,7 @@ exclude_patterns = [
     'test/cpp/jit/upgrader_models/*.ptl',
     'test/cpp/jit/upgrader_models/*.ptl.ff',
     '.lintrunner.toml',
+    'cmake/External/nccl.patch',
 ]
 command = [
     'python3',

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -355,6 +355,8 @@ option(USE_PER_OPERATOR_HEADERS "Whether ATen should generate separate headers f
 cmake_dependent_option(
     BUILD_LAZY_TS_BACKEND "Build the lazy Torchscript backend, not compatible with mobile builds" ON
     "NOT INTERN_BUILD_MOBILE" OFF)
+cmake_dependent_option(
+    BUILD_FUNCTORCH "Build Functorch" ON "BUILD_PYTHON" OFF)
 
 
 if(USE_CCACHE)
@@ -622,6 +624,7 @@ if(INTERN_BUILD_MOBILE)
     set(INTERN_DISABLE_AUTOGRAD ON)
   endif()
   set(BUILD_PYTHON OFF)
+  set(BUILD_FUNCTORCH OFF)
   set(BUILD_CAFFE2_OPS OFF)
   set(USE_DISTRIBUTED OFF)
   set(NO_API ON)
@@ -1175,3 +1178,7 @@ caffe2_print_configuration_summary()
 if(USE_DEPLOY)
   add_subdirectory(torch/csrc/deploy)
 endif()
+
+if(BUILD_FUNCTORCH)
+  add_subdirectory(functorch)
+endif()
diff --git a/MANIFEST.in b/MANIFEST.in
@@ -26,5 +26,6 @@ recursive-include benchmarks *.*
 recursive-include scripts *.*
 recursive-include mypy_plugins *.*
 recursive-include modules *.*
+recursive-include functorch *.*
 prune */__pycache__
 global-exclude *.o *.so *.dylib *.a .git *.pyc *.swp
diff --git a/aten/src/ATen/FunctionalInverses.cpp b/aten/src/ATen/FunctionalInverses.cpp
@@ -228,6 +228,11 @@ Tensor FunctionalInverses::transpose_copy_int_inverse(const Tensor& base, const
     }
 }
 
+Tensor FunctionalInverses::_nested_view_from_buffer_copy_inverse(const Tensor& base, const Tensor& mutated_view, bool reapply_views, const Tensor& nested_size_tensor, const Tensor& nested_stride_tensor, IntArrayRef offsets) {
+    TORCH_INTERNAL_ASSERT(false, "Attempted to call _nested_view_from_buffer() during the functionalization pass. For now, nested tensors aren't supported during functionalization");
+    return Tensor();
+}
+
 Tensor FunctionalInverses::unsqueeze_copy_inverse(const Tensor& base, const Tensor& mutated_view, bool reapply_views, int64_t dim) {
     if (reapply_views) {
       return at::squeeze(mutated_view, dim);

diff --git a/aten/src/ATen/InferSize.h b/aten/src/ATen/InferSize.h
@@ -2,6 +2,8 @@
 
 #include <ATen/DimVector.h>
 #include <c10/core/ScalarType.h>
+#include <c10/core/SymIntArrayRef.h>
+#include <c10/util/DimVector.h>
 #include <c10/util/Optional.h>
 #include <sstream>
 #include <vector>
@@ -14,9 +16,13 @@ namespace at {
 // templated to handle std::vector<int64_t> and DimVector use cases, see
 // below
 //
-template <typename ResultVec>
-inline void infer_size_impl(IntArrayRef shape, int64_t numel, ResultVec& res) {
-  int64_t newsize = 1;
+template <typename InputArrayRef, typename NumelType, typename ResultVec>
+inline void infer_size_impl(
+    InputArrayRef shape,
+    NumelType numel,
+    ResultVec& res) {
+  NumelType newsize = 1;
+  // N.B. this is an index, not a sym dim!
   auto infer_dim = c10::optional<int64_t>();
   for (int64_t dim = 0, ndim = shape.size(); dim != ndim; dim++) {
     if (shape[dim] == -1) {
@@ -69,4 +75,13 @@ inline at::DimVector infer_size_dv(IntArrayRef shape, int64_t numel) {
   return res;
 }
 
+inline at::SymDimVector infer_size_dv(
+    c10::SymIntArrayRef shape,
+    c10::SymInt numel) {
+  auto res = at::SymDimVector(shape);
+  infer_size_impl<c10::SymIntArrayRef, c10::SymInt, at::SymDimVector>(
+      shape, numel, res);
+  return res;
+}
+
 } // namespace at
diff --git a/aten/src/ATen/NestedTensorImpl.cpp b/aten/src/ATen/NestedTensorImpl.cpp
@@ -4,6 +4,7 @@
 #include <ATen/core/op_registration/op_registration.h>
 #include <ATen/NestedTensorImpl.h>
 #include <c10/core/DispatchKey.h>
+#include <c10/core/DispatchKeySet.h>
 #include <c10/util/Exception.h>
 #include <c10/core/TensorImpl.h>
 
@@ -25,6 +26,46 @@ inline void validate_nested_tensor_metadata(
       (size_dim == 0 && (int64_t)offsets.empty()) ||
       (size_dim == 2 && nested_sizes.size(0) == (int64_t)offsets.size()));
 }
+
+/**
+ * Generates a nested key_set from a non-nested tensor.
+ *
+ * When creating a nested tensor from a non-nested tensor
+ * We want to maintain the same keyset as the buffer but
+ * swap non nested keys for nested ones
+ *
+ * @return Appropriate key set for nested tensor
+ */
+inline c10::DispatchKeySet generate_nested_key_set_from_buffer(
+    const at::Tensor& buffer) {
+  auto nested_key_set = buffer.key_set();
+  const bool has_autograd = nested_key_set.has_any(c10::autograd_dispatch_keyset);
+  // Remove non_nested tensor specific keys
+  nested_key_set = nested_key_set -
+      c10::DispatchKeySet{c10::DispatchKey::Dense, c10::DispatchKey::Autograd};
+
+  // Add nested tensor specific keys
+  nested_key_set =
+      nested_key_set | c10::DispatchKeySet{c10::DispatchKey::NestedTensor};
+  nested_key_set =
+      has_autograd ? nested_key_set | c10::autograd_nested : nested_key_set;
+  return nested_key_set;
+}
+
+/**
+ * Generates a the correct view keyset.
+ *
+ * When creating a nested tensor view of base
+ * The appropriate keyset will be dependent on the nested
+ * status of the base
+ *
+ * @return Appropriate key set for nested tensor
+ */
+c10::DispatchKeySet get_view_key_set(const at::Tensor& base) {
+  return base.is_nested() ? base.key_set()
+                          : generate_nested_key_set_from_buffer(base);
+}
+
 } // namespace
 namespace at {
 namespace native {
@@ -119,19 +160,6 @@ inline std::vector<int64_t> construct_offsets(const at::Tensor& sizes) {
   return offsets;
 }
 
-// [Note: Nested Tensor Autograd] The Nested Tensor key is a functionality
-// key and therefore getAutogradRelatedKeySetFromBackend will return the
-// wrong autograd key. For this specific impl we make sure to register the
-// correct Autograd key which is AutogradNestedTensor
-c10::DispatchKeySet generate_nested_key_set(at::Tensor buffer) {
-  c10::DispatchKeySet key_set =
-      c10::DispatchKeySet(DispatchKey::NestedTensor) | c10::DispatchKeySet{buffer.key_set().highestBackendKey()};
-
-  // Add AutogradNestedTensor specific keys
-  key_set = key_set | inplace_or_view_ks | autograd_nested;
-  return key_set;
-}
-
 NestedTensorImpl::NestedTensorImpl(
     Storage storage,
     c10::DispatchKeySet key_set,
@@ -164,7 +192,7 @@ NestedTensorImpl::NestedTensorImpl(
     std::vector<int64_t>&& offsets)
     : NestedTensorImpl(
           buffer.storage(),
-          generate_nested_key_set(buffer),
+          generate_nested_key_set_from_buffer(buffer),
           buffer.dtype(),
           nested_size_tensor,
           nested_stride_tensor,
@@ -195,12 +223,11 @@ NestedTensorImpl::NestedTensorImpl(
     at::Tensor nested_size_tensor,
     at::Tensor nested_stride_tensor,
     std::vector<int64_t>&& offsets)
-    : TensorImpl(impl_type, Storage(base_tensor.storage()), base_tensor.key_set(), base_tensor.dtype()),
+    : TensorImpl(impl_type, Storage(base_tensor.storage()), get_view_key_set(base_tensor), base_tensor.dtype()),
       nested_size_tensor_(std::move(nested_size_tensor)),
       nested_stride_tensor_(std::move(nested_stride_tensor)),
       offsets_(std::move(offsets)),
       opt_sizes_(construct_opt_sizes(nested_size_tensor_)) {
-  TORCH_INTERNAL_ASSERT(base_tensor.is_nested());
   validate_nested_tensor_metadata(nested_size_tensor_, nested_stride_tensor_, offsets_);
   refresh_dim();
   set_custom_sizes_strides(c10::TensorImpl::SizesStridesPolicy::CustomSizes);

diff --git a/aten/src/ATen/NestedTensorImpl.h b/aten/src/ATen/NestedTensorImpl.h
@@ -167,7 +167,7 @@ struct TORCH_API NestedTensorImpl : public c10::TensorImpl {
    * is generated and redispatched to a non-nested kernel this function
    * generates the key set used by that buffer tensor
    *
-   * @return A newly constructed view tensor
+   * @return Appropriate key set for non-nested tensor
    */
   inline c10::DispatchKeySet generate_buffer_key_set() const {
     auto buffer_key_set = this->key_set();
@@ -184,6 +184,7 @@ struct TORCH_API NestedTensorImpl : public c10::TensorImpl {
     buffer_key_set = Autograd
         ? c10::DispatchKeySet{c10::DispatchKey::Autograd} | buffer_key_set
         : buffer_key_set;
+
     return buffer_key_set;
   }
 };
Original file line number	Diff line number	Diff line change
		@@ -1 +1 @@
		ffd056dc1510bdfecafb689ed87601055694f3e6
		138eb96c36c6ba5a14e65d28fad3cbd5a6e6dcfc
Original file line number	Diff line number	Diff line change
		@@ -1 +1 @@
		a89b1957a62e2f68f001d5d60268743edbe164d8
		ebb68f337de8f66df5bb83f8266f3a9c7020094e