Revert "Require less alignment for attn bias (#114173) (#114837)"

This reverts commit 59656491f3b1da809312942872cce010337504b0.

aten/src/ATen/ExpandUtils.h

@@ -187,17 +187,18 @@ expand_inplace(
 // See NOTE [ ExpandUtils Borrowing ] above for `MaybeOwned` explanation.
 inline std::tuple<c10::MaybeOwned<Tensor>, c10::MaybeOwned<Tensor>>
 expand_outplace(const Tensor& to_expand1, const Tensor& to_expand2) {
-  if (to_expand1.sizes().equals(to_expand2.sizes())) {
+  auto s1 = to_expand1.sym_sizes();
+  auto s2 = to_expand2.sym_sizes();
+  if (s1.equals(s2)) {
     return std::make_tuple(
         c10::MaybeOwned<Tensor>::borrowed(to_expand1),
         c10::MaybeOwned<Tensor>::borrowed(to_expand2));
   }
 
-  auto expanded_size =
-      infer_size_dimvector(to_expand1.sizes(), to_expand2.sizes());
+  auto expanded_size = infer_size_symdimvector(s1, s2);
   return std::make_tuple(
-      c10::MaybeOwned<Tensor>::owned(to_expand1.expand(expanded_size)),
-      c10::MaybeOwned<Tensor>::owned(to_expand2.expand(expanded_size)));
+      c10::MaybeOwned<Tensor>::owned(to_expand1.expand_symint(expanded_size)),
+      c10::MaybeOwned<Tensor>::owned(to_expand2.expand_symint(expanded_size)));
 }
 
 inline std::tuple<c10::MaybeOwned<Tensor>, c10::MaybeOwned<Tensor>>

aten/src/ATen/mps/MPSStream.mm

@@ -147,9 +147,9 @@ void MPSStream::addCompletedHandler(MTLCommandBufferHandler block) {
 }
 
 void MPSStream::fill(id<MTLBuffer> buffer, uint8_t value, size_t length, size_t offset, SyncType syncType) {
-  TORCH_INTERNAL_ASSERT(length >= offset);
-  if (length == 0)
+  if (length == 0) {
     return;
+  }
   dispatch_sync(_serialQueue, ^() {
     @autoreleasepool {
       endKernelCoalescing();

aten/src/ATen/native/Distance.cpp

@@ -308,16 +308,18 @@ Tensor cosine_similarity(const Tensor& x1_, const Tensor& x2_, int64_t dim, doub
   // We accept integral types (and bools lol) but vector_norm does not
   auto x1_is_int = c10::isIntegralType(x1_.scalar_type(), /*încludeBool=*/true);
   auto x2_is_int = c10::isIntegralType(x2_.scalar_type(), /*încludeBool=*/true);
-  auto x1 = x1_is_int ? x1_.to(commonDtype) : x1_;
-  auto x2 = x2_is_int ? x2_.to(commonDtype) : x2_;
+  auto x1_t = x1_is_int ? x1_.to(commonDtype) : x1_;
+  auto x2_t = x2_is_int ? x2_.to(commonDtype) : x2_;
+  c10::MaybeOwned<Tensor> x1, x2;
+  std::tie(x1, x2) = expand_outplace(x1_t, x2_t);
 
 
   // We want to divide each tensor by its norm first, as it's more numerically stable.
   // This keeps the result between -1.0 and 1.0
   // We clone them, as we're going to modify them in-place
   // This allows the gradients to propagate propertly all the way to x1 and x2
-  auto x1_norm = at::linalg_vector_norm(x1, 2, /*dim=*/dim, /*keepdim=*/true).clone();
-  auto x2_norm = at::linalg_vector_norm(x2, 2, /*dim=*/dim, /*keepdim=*/true).clone();
+  auto x1_norm = at::linalg_vector_norm(*x1, 2, /*dim=*/dim, /*keepdim=*/true).clone();
+  auto x2_norm = at::linalg_vector_norm(*x2, 2, /*dim=*/dim, /*keepdim=*/true).clone();
 
   {
     at::NoGradGuard guard;
@@ -325,7 +327,7 @@ Tensor cosine_similarity(const Tensor& x1_, const Tensor& x2_, int64_t dim, doub
     x2_norm.clamp_min_(eps);
   }
 
-  return ((x1 / x1_norm) * (x2 / x2_norm)).sum(dim);
+  return ((*x1 / x1_norm) * (*x2 / x2_norm)).sum(dim);
 }
 
 }}  // namespace at::native

aten/src/ATen/native/LinearAlgebra.cpp

@@ -1483,12 +1483,14 @@ static void addmm_impl_cpu_(
   // it is faster to call oneDNN matrix multiplication primitive with RHS*LHS
   // that will call then into Arm® Compute Library (ACL) GEMM kernel and also
   // additionally have support for running kernel with BF16 instructions
-  bool apply_heur = apply_mkldnn_matmul_heur(b.sizes()[0], b.sizes()[1], a.sizes()[1]);
-  if (apply_heur && transpose_a && !transpose_b && result.scalar_type() == at::ScalarType::Float) {
-      mkldnn_matmul(b, a, c, beta.to<float>(), alpha.to<float>());
-      // We have dispatched to ACL GEMM for single precision float
-      // so do not need to dispatch to BLAS GEMM below
-      dispatched = true;
+  if (transpose_c) {
+    bool apply_heur = apply_mkldnn_matmul_heur(b.sizes()[0], b.sizes()[1], a.sizes()[1]);
+    if (apply_heur && transpose_a && !transpose_b && result.scalar_type() == at::ScalarType::Float) {
+        mkldnn_matmul(b, a, c, beta.to<float>(), alpha.to<float>());
+        // We have dispatched to ACL GEMM for single precision float
+        // so do not need to dispatch to BLAS GEMM below
+        dispatched = true;
+    }
   }
 #endif
 

aten/src/ATen/native/cpu/BinaryOpsKernel.cpp

@@ -101,7 +101,7 @@ void mul_kernel(TensorIteratorBase& iter) {
           using comp_t = c10::complex<float>;
           return comp_t{a} * comp_t{b};
         });
-  } else if (iter.is_scalar(2) && at::isReducedFloatingType(dtype)) {
+  } else if (iter.is_scalar(2) && iter.data_ptr(2) != nullptr && at::isReducedFloatingType(dtype)) {
     AT_DISPATCH_REDUCED_FLOATING_TYPES(dtype, "mul_cpu_reduced_float", [&]() {
       using opmath_t = at::opmath_type<scalar_t>;
       opmath_t b = iter.original_scalar_value<opmath_t>(2);
@@ -125,7 +125,7 @@ void mul_kernel(TensorIteratorBase& iter) {
 
 void div_true_kernel(TensorIteratorBase& iter) {
   const auto dtype = iter.common_dtype();
-  if (iter.is_scalar(2) && at::isReducedFloatingType(dtype)) {
+  if (iter.is_scalar(2) && iter.data_ptr(2) != nullptr && at::isReducedFloatingType(dtype)) {
     AT_DISPATCH_REDUCED_FLOATING_TYPES(dtype, "div_cpu_reduced_float", [&]() {
       using opmath_t = at::opmath_type<scalar_t>;
       opmath_t b = iter.original_scalar_value<opmath_t>(2);
@@ -162,19 +162,28 @@ void div_trunc_kernel(TensorIteratorBase& iter) {
         return a / b;
       });
     });
-  } else if (iter.is_scalar(2) && at::isReducedFloatingType(dtype)) {
-    AT_DISPATCH_REDUCED_FLOATING_TYPES(dtype, "div_trunc_cpu_reduced_float", [&]() {
-      using opmath_t = at::opmath_type<scalar_t>;
-      opmath_t b = iter.original_scalar_value<opmath_t>(2);
-      iter.remove_operand(2);
-      cpu_kernel_vec(iter,
-        [=](scalar_t a) __ubsan_ignore_float_divide_by_zero__ -> scalar_t {
-          return std::trunc(static_cast<opmath_t>(a) / b);
-        },
-        [=](Vectorized<scalar_t> a) {
-          return binary_op_scalar(a, b, [](const Vectorized<opmath_t>& x, const Vectorized<opmath_t>& y) { return (x / y).trunc(); });
+  } else if (iter.is_scalar(2) && iter.data_ptr(2) != nullptr && at::isReducedFloatingType(dtype)) {
+    AT_DISPATCH_REDUCED_FLOATING_TYPES(
+        dtype, "div_trunc_cpu_reduced_float", [&]() {
+          using opmath_t = at::opmath_type<scalar_t>;
+          opmath_t b = iter.original_scalar_value<opmath_t>(2);
+          iter.remove_operand(2);
+          cpu_kernel_vec(
+              iter,
+              [=](scalar_t a)
+                  __ubsan_ignore_float_divide_by_zero__ -> scalar_t {
+                    return std::trunc(static_cast<opmath_t>(a) / b);
+                  },
+              [=](Vectorized<scalar_t> a) {
+                return binary_op_scalar(
+                    a,
+                    b,
+                    [](const Vectorized<opmath_t>& x,
+                       const Vectorized<opmath_t>& y) {
+                      return (x / y).trunc();
+                    });
+              });
         });
-    });
   } else {
     AT_DISPATCH_FLOATING_TYPES_AND2(kBFloat16, kHalf, dtype, "div_trunc_cpu", [&]() {
       cpu_kernel_vec(iter,
@@ -223,20 +232,25 @@ void div_floor_kernel(TensorIteratorBase& iter) {
     });
   } else {
     // See NOTE: [Floor Division in Python]
-    if (iter.is_scalar(2) && at::isReducedFloatingType(dtype)) {
-      AT_DISPATCH_REDUCED_FLOATING_TYPES(dtype, "div_floor_cpu_reduced_float", [&]() {
-        using opmath_t = at::opmath_type<scalar_t>;
-        opmath_t b = iter.original_scalar_value<opmath_t>(2);
-        iter.remove_operand(2);
-        using vec_t = Vectorized<opmath_t>;
-        cpu_kernel_vec(iter,
-          [=](scalar_t a) -> scalar_t {
-            return div_floor_floating(static_cast<opmath_t>(a), b);
-          },
-          [=](Vectorized<scalar_t> a) {
-            return binary_op_scalar(a, b, [](const vec_t& x, const vec_t& y) { return div_floor_floating_vec(x, y); });
+    if (iter.is_scalar(2) && iter.data_ptr(2) != nullptr && at::isReducedFloatingType(dtype)) {
+      AT_DISPATCH_REDUCED_FLOATING_TYPES(
+          dtype, "div_floor_cpu_reduced_float", [&]() {
+            using opmath_t = at::opmath_type<scalar_t>;
+            opmath_t b = iter.original_scalar_value<opmath_t>(2);
+            iter.remove_operand(2);
+            using vec_t = Vectorized<opmath_t>;
+            cpu_kernel_vec(
+                iter,
+                [=](scalar_t a) -> scalar_t {
+                  return div_floor_floating(static_cast<opmath_t>(a), b);
+                },
+                [=](Vectorized<scalar_t> a) {
+                  return binary_op_scalar(
+                      a, b, [](const vec_t& x, const vec_t& y) {
+                        return div_floor_floating_vec(x, y);
+                      });
+                });
           });
-      });
     } else {
       AT_DISPATCH_FLOATING_TYPES_AND2(kBFloat16, kHalf, dtype, "div_floor_cpu", [&]() {
         using vec_t = Vectorized<scalar_t>;

aten/src/ATen/native/mps/operations/ConstantOps.mm

@@ -72,7 +72,7 @@ static bool fill_mps_tensor_(Tensor& self, uint8_t value) {
   if (self.is_contiguous()) {
     MPSStream* stream = getCurrentMPSStream();
     auto storage_byte_offset = self.storage_offset() * self.itemsize();
-    stream->fill(mps::getMTLBufferStorage(self), 0, self.storage().nbytes(), storage_byte_offset);
+    stream->fill(mps::getMTLBufferStorage(self), value, self.nbytes(), storage_byte_offset);
     return true;
   }
   return false;

aten/src/ATen/native/mps/operations/Convolution.mm

@@ -445,7 +445,7 @@ static Tensor mps_convolution_backward_weights(IntArrayRef weight_size,
     string key = "mps_convolution_backward_weights:" + to_string(stride[0]) + ":" + to_string(stride[1]) + ":" +
         to_string(dilation[0]) + ":" + to_string(dilation[1]) + ":" + to_string(padding[0]) + ":" +
         to_string(padding[1]) + ":" + to_string(groups) + ":" + mem_format_key +
-        getTensorsStringKey({grad_output_t, input_t}) + ":" + string([ns_shape_key UTF8String]);
+        getTensorsStringKey({grad_output_t, input_t, grad_weight_t}) + ":" + string([ns_shape_key UTF8String]);
 
     auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
       MPSGraphConvolution2DOpDescriptor* conv2dDescriptor_ = [[MPSGraphConvolution2DOpDescriptor new] autorelease];

caffe2/utils/threadpool/ThreadPool.cc

@@ -41,8 +41,13 @@ namespace {
 }
 
 size_t getDefaultNumThreads() {
-  CAFFE_ENFORCE(cpuinfo_initialize(), "cpuinfo initialization failed");
-  int numThreads = cpuinfo_get_processors_count();
+  auto numThreads = 1U;
+  if (cpuinfo_initialize()) {
+    numThreads = std::max(cpuinfo_get_processors_count(), 1U);
+  } else {
+    LOG(WARNING) << "cpuinfo initialization failed";
+    numThreads = std::max(std::thread::hardware_concurrency(), 1U);
+  }
 
   bool applyCap = false;
 #if defined(C10_ANDROID)
@@ -109,7 +114,7 @@ size_t getDefaultNumThreads() {
    * detect if we are running under tsan, for now capping the default
    * threadcount to the tsan limit unconditionally.
    */
-  int tsanThreadLimit = 63;
+  auto tsanThreadLimit = 63U;
   numThreads = std::min(numThreads, tsanThreadLimit);
 
   return numThreads;

setup.py

@@ -1323,6 +1323,7 @@ def main():
         "include/torch/csrc/lazy/ts_backend/*.h",
         "include/pybind11/*.h",
         "include/pybind11/detail/*.h",
+        "include/pybind11/eigen/*.h",
         "include/TH/*.h*",
         "include/TH/generic/*.h*",
         "include/THC/*.cuh",

test/test_foreach.py

@@ -516,16 +516,20 @@ class TestForeach(TestCase):
                 sum(ref((ref_tensors,), ord=ord)).backward()
                 self.assertEqual([t.grad for t in tensors], [t.grad for t in ref_tensors])
 
-    @dtypes(*all_types_and_complex_and(torch.half, torch.bfloat16, torch.bool))
+    @dtypes(*all_types_and_complex_and(torch.half, torch.bfloat16))
     def test_add_scalar_with_empty_list_and_empty_tensor(self, device, dtype):
         # TODO: enable empty list case
-        for tensors in [[torch.randn([0])]]:
+        for tensors in [[torch.randn([0], device=device, dtype=dtype)],
+                        [torch.empty_strided((0, 1), (0, 0), dtype=dtype, device=device)]]:
             res = torch._foreach_add(tensors, 1)
             self.assertEqual(res, tensors)
 
             torch._foreach_add_(tensors, 1)
             self.assertEqual(res, tensors)
 
+            # Regression test for https://github.com/pytorch/pytorch/issues/113156
+            torch._foreach_mul_(tensors, 1)
+
     @ops(
         filter(lambda op: not op.has_no_out_of_place, foreach_binary_op_db),
         dtypes=OpDTypes.supported,

test/test_mps.py

@@ -1197,17 +1197,22 @@ class TestMPS(TestCaseMPS):
         tensor_cpu = tensor_0[:][1].fill_(val)
 
         self.assertEqual(tensor_mps, tensor_cpu)
+        self.assertEqual(tensor, tensor_0)
 
         shape = [1, 10]
         val = 0.0
         tensor = torch.ones(shape, device="mps")
         val_tensor_mps = torch.tensor(val, device="mps")
         tensor_mps = tensor[:, 9].fill_(val_tensor_mps)
+        # Regression test for https://github.com/pytorch/pytorch/issues/114692
+        tensor[:, 5].fill_(val_tensor_mps)
         tensor_0 = torch.ones(shape, device="cpu")
         val_tensor_cpu = torch.tensor(val, device="cpu")
         tensor_cpu = tensor_0[:, 9].fill_(val_tensor_cpu)
+        tensor_0[:, 5].fill_(val_tensor_cpu)
 
-        self.assertEqual(tensor_mps, tensor_cpu)
+        self.assertEqual(tensor_mps.to(device="cpu"), tensor_cpu)
+        self.assertEqual(tensor.to(device="cpu"), tensor_0)
 
     def test_cdist_large(self, device="mps"):
         for cm in ['use_mm_for_euclid_dist_if_necessary', 'use_mm_for_euclid_dist', 'donot_use_mm_for_euclid_dist']:
@@ -7982,6 +7987,18 @@ class TestNNMPS(NNTestCase):
         actual = F.conv2d(x, y, padding='valid')
         self.assertEqual(expect.to('cpu'), actual.to('cpu'))
 
+    def test_conv2d_backward_collision(self):
+        # Test for https://github.com/pytorch/pytorch/issues/112998
+        x = torch.rand(1, 1, 10, 10, device="mps", requires_grad=True)
+        m1 = nn.Conv2d(1, 1, 3, stride=2, padding=1).to("mps")
+        m2 = nn.Conv2d(1, 1, 4, stride=2, padding=1).to("mps")
+        y1, y2 = m1(x), m2(x)
+        self.assertEqual(y1.shape, y2.shape)
+        y1.sum().backward()
+        # This used to crash with MPSNDArrayConvolutionA14.mm:4352: failed assertion
+        y2.sum().backward()
+
+
     def test_gemm_permute_transpose(self):
         batch_size = 32
         n = 20

test/test_nn.py

@@ -5609,6 +5609,18 @@ tensor(..., device='meta', size=(1,), requires_grad=True)""")
         out = F.cosine_similarity(input.to(torch.int8), input, dim=-1)
         self.assertEqual(out, 1.)
 
+        # Check broadcasting #109333
+        a = torch.ones(2, 3, dtype=torch.float)
+        b = torch.ones(1, 1, dtype=torch.float)
+        out = F.cosine_similarity(a, b)
+        self.assertEqual(out, torch.ones(2, dtype=torch.float))
+
+        a = torch.ones(2, 3, dtype=torch.float)
+        b = torch.ones(1, dtype=torch.float)
+        out = F.cosine_similarity(a, b)
+        self.assertEqual(out, torch.ones(2, dtype=torch.float))
+
+
     def test_grid_sample_error_checking(self):
         input = torch.empty(1, 1, 2, 2)
         grid = torch.empty(1, 1, 1, 2)

test/test_numpy_interop.py

@@ -472,6 +472,21 @@ class TestNumPyInterop(TestCase):
                 else:
                     self.assertTrue(t == a)
 
+    @onlyCPU
+    def test_empty_tensors_interop(self, device):
+        x = torch.rand((), dtype=torch.float16)
+        y = torch.tensor(np.random.rand(0), dtype=torch.float16)
+        # Same can be achieved by running
+        # y = torch.empty_strided((0,), (0,), dtype=torch.float16)
+
+        # Regression test for https://github.com/pytorch/pytorch/issues/115068
+        self.assertEqual(torch.true_divide(x, y).shape, y.shape)
+        # Regression test for https://github.com/pytorch/pytorch/issues/115066
+        self.assertEqual(torch.mul(x, y).shape, y.shape)
+        # Regression test for https://github.com/pytorch/pytorch/issues/113037
+        self.assertEqual(torch.div(x, y, rounding_mode='floor').shape, y.shape)
+
+
 instantiate_device_type_tests(TestNumPyInterop, globals())
 
 if __name__ == '__main__':

test/test_proxy_tensor.py

@@ -1604,7 +1604,6 @@ symbolic_tensor_failures.update(symbolic_tensor_segfaults)
 
 outplace_symbolic_tensor_failures = {
     xfail('i0', ''),  # aten.i0.default - couldn't find symbolic meta function/decomposition
-    xfail('masked_scatter', ''),  # aten.masked_scatter.default - couldn't find symbolic meta function/decomposition
 }
 
 inplace_symbolic_tensor_failures = {

test/test_torch.py

@@ -9151,6 +9151,13 @@ tensor([[[1.+1.j, 1.+1.j, 1.+1.j,  ..., 1.+1.j, 1.+1.j, 1.+1.j],
         t2 = t[0:0].view(0, 1)
         self.assertEqual(t2.data_ptr(), 0)
 
+    def test_invalid_arg_error_handling(self) -> None:
+        """ Tests that errors from old TH functions are propagated back """
+        for invalid_val in [-1, 2**65]:
+            self.assertRaises(RuntimeError, lambda: torch.set_num_threads(invalid_val))
+            self.assertRaises(RuntimeError, lambda: torch.set_num_interop_threads(invalid_val))
+
+
 # The following block extends TestTorch with negative dim wrapping tests
 # FIXME: replace these with OpInfo sample inputs or systemic OpInfo tests
 # Functions to test negative dimension wrapping

tools/stats/import_test_stats.py

@@ -66,7 +66,7 @@ def fetch_and_cache(
 def get_slow_tests(
     dirpath: str, filename: str = SLOW_TESTS_FILE
 ) -> Optional[Dict[str, float]]:
-    url = "https://ossci-metrics.s3.amazonaws.com/slow-tests.json"
+    url = "https://ossci-metrics.s3.amazonaws.com/slow-tests.json?versionId=iWAOsEqlVH1mfs7w5A3KlyTalvubE4Ru"
     try:
         return fetch_and_cache(dirpath, filename, url, lambda x: x)
     except Exception:
@@ -98,7 +98,7 @@ def get_disabled_tests(
         return disabled_test_from_issues
 
     try:
-        url = "https://ossci-metrics.s3.amazonaws.com/disabled-tests-condensed.json"
+        url = "https://ossci-metrics.s3.amazonaws.com/disabled-tests-condensed.json?versionId=JMUOdxgUAeI4yXhzc.dJlCuxVrsfkZTj"
         return fetch_and_cache(dirpath, filename, url, process_disabled_test)
     except Exception:
         print("Couldn't download test skip set, leaving all tests enabled...")

torch/autograd/profiler_util.py

@@ -253,7 +253,7 @@ class EventList(list):
                         '"pid": "CPU functions", '
                         f'"id": {next_id}, '
                         f'"cat": "cpu_to_{device_name}", '
-                        '"args": {{}}}}, '
+                        '"args": {}}, '
                     )
                     # Note: use torch.profiler to get device kernel trace
                     next_id += 1

torch/csrc/Module.cpp

@@ -240,6 +240,7 @@ static PyObject* THPModule_getNumThreads(PyObject* module, PyObject* noargs) {
 }
 
 static PyObject* THPModule_setNumThreads(PyObject* module, PyObject* arg) {
+  HANDLE_TH_ERRORS
   THPUtils_assert(
       THPUtils_checkLong(arg),
       "set_num_threads expects an int, "
@@ -249,6 +250,7 @@ static PyObject* THPModule_setNumThreads(PyObject* module, PyObject* arg) {
   THPUtils_assert(nthreads > 0, "set_num_threads expects a positive integer");
   at::set_num_threads(nthreads);
   Py_RETURN_NONE;
+  END_HANDLE_TH_ERRORS
 }
 
 static PyObject* THPModule_getNumInteropThreads(
@@ -260,6 +262,7 @@ static PyObject* THPModule_getNumInteropThreads(
 static PyObject* THPModule_setNumInteropThreads(
     PyObject* module,
     PyObject* arg) {
+  HANDLE_TH_ERRORS
   THPUtils_assert(
       THPUtils_checkLong(arg),
       "set_num_interop_threads expects an int, "
@@ -270,6 +273,7 @@ static PyObject* THPModule_setNumInteropThreads(
       nthreads > 0, "set_num_interop_threads expects a positive integer");
   at::set_num_interop_threads(nthreads);
   Py_RETURN_NONE;
+  END_HANDLE_TH_ERRORS
 }
 
 PyObject* THPModule_setDefaultTensorType(PyObject* _unused, PyObject* type) {