2023-01-12 nightly release (5766764)

pytorch · ezyang · Sep 11, 2022 · Sep 12, 2022 · Sep 13, 2022 · Sep 14, 2022
commit a0c46f5ab85792df142fc8a4454ad0fa4a50d2d6
@@ -0,0 +1,2 @@
+typing_extensions
+jinja2
@@ -22,7 +22,7 @@ couple of utility templates used to discern common utilities that can be used am
 
 You will need `jinja2` in order to regenerate the workflow files which can be installed using:
 ```bash
-pip install -r .github/requirements.txt
+pip install -r .github/requirements/regenerate-requirements.txt
 ```
 
 Workflows can be generated / regenerated using the following command:
@@ -56,3 +56,10 @@ Exceptions to AWS for self hosted:
 New runner types can be added by committing changes to `.github/scale-config.yml`. Example: https://github.com/pytorch/pytorch/pull/70474
 
 > NOTE: New runner types can only be used once the changes to `.github/scale-config.yml` have made their way into the default branch
+
+### Testing [pytorch/builder](https://github.com/pytorch/builder) changes
+
+In order to test changes to the builder scripts:
+
+1. Specify your builder PR's branch and repo as `builder_repo` and  `builder_branch` in [`.github/templates/common.yml.j2`](https://github.com/pytorch/pytorch/blob/32356aaee6a77e0ae424435a7e9da3d99e7a4ca5/.github/templates/common.yml.j2#LL10C26-L10C32). 2. Regenerate workflow files with `.github/regenerate.sh` (see above).
+3. Submit fake PR to PyTorch. If changing binaries build, add an appropriate label like `ciflow/binaries` to trigger the builds.
@@ -7,6 +7,7 @@
 
 # NOTE: If testing pytorch/builder changes you can change this variable to change what pytorch/builder reference
 #       the binary builds will check out
+{%- set builder_repo = "pytorch/builder" -%}
 {%- set builder_branch = "main" -%}
 
 {%- macro concurrency(build_environment) -%}

@@ -88,7 +88,7 @@ jobs:
           name: !{{ config["build_name"] }}
           path: "${{ runner.temp }}/artifacts/"
       !{{ common.checkout(deep_clone=False, directory="pytorch") }}
-      !{{ common.checkout(deep_clone=False, directory="builder", repository="pytorch/builder", branch=common.builder_branch) }}
+      !{{ common.checkout(deep_clone=False, directory="builder", repository=common.builder_repo, branch=common.builder_branch) }}
       - name: ROCm set GPU_FLAG
         run: |
           echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd --device=/dev/dri --group-add video --group-add daemon" >> "${GITHUB_ENV}"

@@ -75,7 +75,7 @@ jobs:
           echo "${RUNNER_TEMP}/anaconda/bin" >> "${GITHUB_PATH}"
           echo "DEVELOPER_DIR=/Applications/Xcode_13.3.1.app/Contents/Developer" >> "${GITHUB_ENV}"
       !{{ common.checkout(deep_clone=False, directory="pytorch") }}
-      !{{ common.checkout(deep_clone=False, directory="builder", repository="pytorch/builder", branch=common.builder_branch) }}
+      !{{ common.checkout(deep_clone=False, directory="builder", repository=common.builder_repo, branch=common.builder_branch) }}
       - name: Install sccache (only for non-forked PRs, and pushes to trunk)
         uses: nick-fields/retry@v2.8.2
         if: ${{ github.event_name == 'push' || github.event.pull_request.head.repo.full_name == github.repository }}

@@ -63,7 +63,7 @@ jobs:
       !{{ common.setup_ec2_windows() }}
       !{{ set_runner_specific_vars() }}
       !{{ common.checkout(deep_clone=False, directory="pytorch") }}
-      !{{ common.checkout(deep_clone=False, directory="builder", repository="pytorch/builder", branch=common.builder_branch) }}
+      !{{ common.checkout(deep_clone=False, directory="builder", repository=common.builder_repo, branch=common.builder_branch) }}
       - name: Populate binary env
         shell: bash
         run: |
@@ -99,7 +99,7 @@ jobs:
           name: !{{ config["build_name"] }}
           path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}"
       !{{ common.checkout(deep_clone=False, directory="pytorch") }}
-      !{{ common.checkout(deep_clone=False, directory="builder", repository="pytorch/builder", branch=common.builder_branch) }}
+      !{{ common.checkout(deep_clone=False, directory="builder", repository=common.builder_repo, branch=common.builder_branch) }}
       - name: Populate binary env
         shell: bash
         run: |

@@ -230,7 +230,6 @@ jobs:
       test-matrix: ${{ needs.linux-bionic-py3_7-clang8-xla-build.outputs.test-matrix }}
 
   win-vs2019-cpu-py3-build:
-    if: false
     name: win-vs2019-cpu-py3
     uses: ./.github/workflows/_win-build.yml
     with:
@@ -244,7 +243,6 @@ jobs:
         ]}
 
   win-vs2019-cpu-py3-test:
-    if: false
     name: win-vs2019-cpu-py3
     uses: ./.github/workflows/_win-test.yml
     needs: win-vs2019-cpu-py3-build

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -529,6 +529,16 @@ if(MSVC)
     string(APPEND ${flag_var} " /ignore:4049 /ignore:4217 /ignore:4099")
   endforeach(flag_var)
 
+  foreach(flag_var
+      CMAKE_SHARED_LINKER_FLAGS)
+    # https://github.com/pytorch/pytorch/issues/91933: Don't set the manifest filename
+    # explicitly helps fix the linker error when linking torch_python.dll. The manifest
+    # file would still be there in the correct format torch_python.dll.manifest
+    if(${flag_var} MATCHES "/MANIFESTFILE:.*\\.manifest")
+      string(REGEX REPLACE "/MANIFESTFILE:.*\\.manifest" "" ${flag_var} "${${flag_var}}")
+    endif()
+  endforeach(flag_var)
+
   # Try harder
   string(APPEND CMAKE_CUDA_FLAGS " -Xcompiler /w -w")
 

diff --git a/aten/src/ATen/PythonTorchFunctionTLS.cpp b/aten/src/ATen/PythonTorchFunctionTLS.cpp
@@ -26,12 +26,12 @@ int64_t PythonTorchFunctionTLS::stack_len() {
   return pythonTorchFunctionState.stack_.size();
 }
 
-void PythonTorchFunctionTLS::set_disabled(bool disabled) {
-  pythonTorchFunctionState.disabled_ = disabled;
+void PythonTorchFunctionTLS::set_disabled_state(TorchFunctionDisabledState disabled_state) {
+  pythonTorchFunctionState.disabled_state_ = disabled_state;
 }
 
-bool PythonTorchFunctionTLS::is_disabled() {
-  return pythonTorchFunctionState.disabled_;
+TorchFunctionDisabledState PythonTorchFunctionTLS::get_disabled_state() {
+  return pythonTorchFunctionState.disabled_state_;
 }
 
 void PythonTorchFunctionTLS::set_state(const PythonTorchFunctionTLS& state) {
@@ -43,7 +43,8 @@ const PythonTorchFunctionTLS& PythonTorchFunctionTLS::get_state() {
 }
 
 bool torch_function_mode_enabled() {
-  return PythonTorchFunctionTLS::stack_len() > 0;
+  return PythonTorchFunctionTLS::get_disabled_state() != TorchFunctionDisabledState::ALL_DISABLED &&
+         PythonTorchFunctionTLS::stack_len() > 0;
 }
 
 } // namespace impl

diff --git a/aten/src/ATen/PythonTorchFunctionTLS.h b/aten/src/ATen/PythonTorchFunctionTLS.h
@@ -6,9 +6,11 @@
 namespace at {
 namespace impl {
 
+enum TorchFunctionDisabledState { ENABLED, SUBCLASSES_DISABLED, ALL_DISABLED };
+
 struct TORCH_API PythonTorchFunctionTLS {
-  static void set_disabled(bool);
-  static bool is_disabled();
+  static void set_disabled_state(TorchFunctionDisabledState disabled_state_);
+  static TorchFunctionDisabledState get_disabled_state();
 
   static void push_onto_stack(std::shared_ptr<SafePyObject> mode);
   static const std::shared_ptr<SafePyObject> pop_stack();
@@ -20,11 +22,11 @@ struct TORCH_API PythonTorchFunctionTLS {
 
  private:
   // The mode TLS is split into
-  //   - disabled_, which says whether or not to disable all torch function
-  //   modes
+  //   - disabled_state, which says which part of torch function are disabled
   //   - stack_, which is a vector of modes representing the stack of user
   //   defined modes
-  bool disabled_;
+  TorchFunctionDisabledState disabled_state_ =
+      TorchFunctionDisabledState::ENABLED;
   std::vector<std::shared_ptr<c10::SafePyObject>> stack_;
 };
 

diff --git a/aten/src/ATen/cpu/vml.h b/aten/src/ATen/cpu/vml.h
@@ -2,6 +2,7 @@
 
 #include <ATen/Config.h>
 #include <ATen/Parallel.h>
+#include <ATen/OpMathType.h>
 #include <ATen/cpu/vec/functional.h>
 #include <ATen/cpu/vec/vec.h>
 #include <c10/util/complex.h>
@@ -34,9 +35,6 @@
 #include <mkl.h>
 #endif
 
-#define DL_RUNTIME_BUG(op, type_)
-#define DL_RUNTIME_BUG_BFLOAT16()
-
 namespace at {
 namespace vml {
 inline namespace CPU_CAPABILITY {
@@ -58,44 +56,11 @@ inline void vrsqrt(scalar_t* out, scalar_t* in, int64_t size) {
 
 // NB: We ignore numerical errors by convention and leave them to the user
 
-// We unfortunately need to duplicate code here to deal with the SSE-AVX
-// transition bug (see [Note SSE-AVX transitions]). As soon as we can expect
-// users to use a version of glibc newer than 2.23 we will be able to ditch
-// this. This duplication is also necessary since not all functions (e.g. rsqrt)
-// might be part of cmath.
-
-// for BFloat16, we need specialize it, the reason is that avx/avx2 and glic=2.23,
-// we can't give DL_RUNTIME_BUG volatile type in x = std::op(x);
-
-#define IMPLEMENT_VML_BUG(op)                                                     \
-  template <typename scalar_t>                                                    \
-  inline void v##op(scalar_t* out, const scalar_t* in, int64_t size) {            \
-    DL_RUNTIME_BUG(op, scalar_t)                                                  \
-    parallel_for(0, size, 2048, [out, in](int64_t begin, int64_t end) {           \
-      map([](const Vectorized<scalar_t>& x) { return x.op(); },                   \
-          out + begin,                                                            \
-          in + begin,                                                             \
-          end - begin);                                                           \
-    });                                                                           \
-  }                                                                               \
-  template <>                                                                     \
-  inline void v##op<c10::BFloat16>(                                               \
-      c10::BFloat16* out, const c10::BFloat16* in, int64_t size) {                \
-    parallel_for(0, size, 2048, [out, in](int64_t begin, int64_t end) {           \
-      DL_RUNTIME_BUG_BFLOAT16()                                                   \
-      using vecscalar_t = vec_scalar_t<c10::BFloat16>;                            \
-      map([](const Vectorized<vecscalar_t>& x) { return x.op(); },                \
-          out + begin,                                                            \
-          in + begin,                                                             \
-          end - begin);                                                           \
-    });                                                                           \
-  }
-
 #define IMPLEMENT_VML(op)                                                         \
   template <typename scalar_t>                                                    \
   inline void v##op(scalar_t* out, const scalar_t* in, int64_t size) {            \
     parallel_for(0, size, 2048, [out, in](int64_t begin, int64_t end) {           \
-      using vecscalar_t = vec_scalar_t<scalar_t>;                                 \
+      using vecscalar_t = at::opmath_type<scalar_t>;                              \
       map([](const Vectorized<vecscalar_t>& x) { return x.op(); },                \
           out + begin,                                                            \
           in + begin,                                                             \
@@ -109,7 +74,7 @@ IMPLEMENT_VML(asin)
 IMPLEMENT_VML(atan)
 IMPLEMENT_VML(ceil)
 IMPLEMENT_VML(cos)
-// IMPLEMENT_VML_BUG(cosh)
+// IMPLEMENT_VML(cosh)
 IMPLEMENT_VML(erf)
 IMPLEMENT_VML(erfc)
 IMPLEMENT_VML(erfinv)
@@ -125,7 +90,7 @@ IMPLEMENT_VML(log1p)
 IMPLEMENT_VML(log2)
 IMPLEMENT_VML(neg)
 IMPLEMENT_VML(sin)
-// IMPLEMENT_VML_BUG(sinh)
+// IMPLEMENT_VML(sinh)
 IMPLEMENT_VML(sqrt)
 IMPLEMENT_VML(round)
 IMPLEMENT_VML(rsqrt)

diff --git a/aten/src/ATen/native/TensorAdvancedIndexing.cpp b/aten/src/ATen/native/TensorAdvancedIndexing.cpp
@@ -492,6 +492,10 @@ DEFINE_DISPATCH(scatter_reduce_stub);
 DEFINE_DISPATCH(scatter_scalar_reduce_stub);
 DEFINE_DISPATCH(scatter_reduce_two_stub);
 
+DEFINE_DISPATCH(scatter_add_expanded_index_stub);
+DEFINE_DISPATCH(scatter_reduce_expanded_index_stub);
+DEFINE_DISPATCH(gather_expanded_index_stub);
+
 static bool all_strides_match(TensorList tensors) {
   TORCH_CHECK(tensors.size() >= 1);
   auto strides = tensors[0].strides();
@@ -1482,7 +1486,11 @@ TORCH_IMPL_FUNC(gather_out)
 (const Tensor& self, int64_t dim, const Tensor& index, bool sparse_grad, const Tensor& result) {
   if (index.numel() == 0) return;
   dim = at::maybe_wrap_dim(dim, self.dim());
-  gather_stub(result.device().type(), result, self, dim, index);
+  if (can_use_expanded_index_path(result, dim, index, self, /*is_scatter_like=*/false)) {
+    gather_expanded_index_stub(result.device().type(), result, self, index);
+  } else {
+    gather_stub(result.device().type(), result, self, dim, index);
+  }
 }
 
 Tensor gather_backward(const Tensor& grad, const Tensor& self, int64_t dim, const Tensor& index, bool sparse_grad) {
@@ -1723,7 +1731,11 @@ TORCH_IMPL_FUNC(scatter_add)
       }
     }
   } else {
-    scatter_add_stub(self.device().type(), mut_out, dim, index, src);
+    if (can_use_expanded_index_path(mut_out, dim, index, src, /*is_scatter_like*/true)) {
+      scatter_add_expanded_index_stub(self.device().type(), mut_out, index, src);
+    } else {
+      scatter_add_stub(self.device().type(), mut_out, dim, index, src);
+    }
   }
 }
 
@@ -1738,13 +1750,26 @@ TORCH_IMPL_FUNC(scatter_reduce_two)
   // See issue https://github.com/pytorch/pytorch/issues/74770
   TORCH_WARN_ONCE("scatter_reduce() is in beta and the API may change at any time.");
 
+  dim = at::maybe_wrap_dim(dim, self.dim());
+
+  if (!self.is_same(out)) {
+    out.copy_(self);
+  }
+
+  const auto op = meta::get_operator_enum(reduce, true);
+
+  if (can_use_expanded_index_path(out, dim, index, src, /*is_scatter_like*/true)) {
+    scatter_reduce_expanded_index_stub(self.device().type(), out, index, src, op, include_self);
+    return;
+  }
+
   scatter_impl</*use_new_options=*/true>(self, dim, index, src, out,
                                          scatter_reduce_two_stub,
                                          scatter_stub,
                                          reduce,
                                          include_self);
 
-  if (meta::get_operator_enum(reduce, true) == SCATTER_GATHER_OP::REDUCE_MEAN) {
+  if (op == SCATTER_GATHER_OP::REDUCE_MEAN) {
     auto ones = at::ones_like(src);
     auto count = include_self ? at::ones_like(out) : at::zeros_like(out);
     count.scatter_add_(dim, index, ones);

diff --git a/aten/src/ATen/native/TensorAdvancedIndexing.h b/aten/src/ATen/native/TensorAdvancedIndexing.h
@@ -5,6 +5,7 @@
 #include <ATen/core/List.h>
 #include <ATen/core/Tensor.h>
 #include <ATen/native/DispatchStub.h>
+#include <ATen/native/cpu/radix_sort.h>
 
 namespace at {
 struct TensorIterator;
@@ -39,4 +40,66 @@ DECLARE_DISPATCH(scatter_reduce_two_fn, scatter_reduce_two_stub);
 
 TORCH_API Tensor& index_out(Tensor& result, const Tensor & self, const c10::List<c10::optional<at::Tensor>>& indices);
 
+// fast paths for GNN usage
+static inline bool can_use_expanded_index_path(
+    const Tensor& self,
+    int64_t dim,
+    const Tensor& index,
+    const Tensor& src,
+    bool is_scatter_like) {
+  if (!self.device().is_cpu()) {
+    return false;
+  }
+
+  const auto st = self.scalar_type();
+  if (!(c10::isFloatingType(st)) || st == ScalarType::Half) {
+    return false;
+  }
+
+  if (!is_radix_sort_available()) {
+    return false;
+  }
+
+  // skip when having empty tensor
+  if (self.numel() == 0 || index.numel() == 0 || src.numel() == 0) {
+    return false;
+  }
+
+  // skip when having scalar tensor
+  if (self.ndimension() == 0 || index.ndimension() == 0 || src.ndimension() == 0) {
+    return false;
+  }
+
+  if (is_scatter_like) {
+    // using `spmm` for scatter would require sorting on index,
+    // this is only perf beneficial when the inner dimension, aka, `channels`
+    // is big enough.
+    constexpr int64_t threshold = 16;
+    if (index.numel() / index.size(0) < threshold) {
+      return false;
+    }
+  }
+
+  // usually the expanded index has stride on the first dimension to be 1,
+  // and strides on other dims to be 0 or 1, e.g.
+  //   shape [108365, 16]; strides [1, 0]
+  //   shape [13264, 1, 7]; strides [1, 1, 0]
+  auto index_strides = index.strides().vec();
+  bool is_index_expanded = index_strides[0] == 1;
+  for (const auto dim : c10::irange(1, index_strides.size())) {
+    if (index_strides[dim] > 1) { is_index_expanded = false; }
+  }
+
+  // index is expanded
+  return dim == 0 && is_index_expanded && src.is_contiguous() && self.is_contiguous();
+}
+
+using scatter_add_expanded_index_fn = void(*)(const Tensor&, const Tensor&, const Tensor&);
+using scatter_reduce_expanded_index_fn = void(*)(const Tensor&, const Tensor&, const Tensor&, const SCATTER_GATHER_OP& reduce, bool);
+using gather_expanded_index_fn = void (*)(const Tensor&, const Tensor&, const Tensor&);
+
+DECLARE_DISPATCH(scatter_add_expanded_index_fn, scatter_add_expanded_index_stub);
+DECLARE_DISPATCH(scatter_reduce_expanded_index_fn, scatter_reduce_expanded_index_stub);
+DECLARE_DISPATCH(gather_expanded_index_fn, gather_expanded_index_stub);
+
 }} // namespace at::native