Support uint{16,32,64} deterministic empty fill and scalar Python binding handling (pytorch#116807)

Signed-off-by: Edward Z. Yang <[email protected]>
Pull Request resolved: pytorch#116807
Approved by: https://github.com/albanD
ghstack dependencies: pytorch#116805, pytorch#116806

Author: ezyang

aten/src/ATen/ScalarOps.cpp

@@ -1,5 +1,6 @@
 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
 #include <ATen/Dispatch.h>
+#include <ATen/Dispatch_v2.h>
 #include <ATen/EmptyTensor.h>
 #include <ATen/ScalarOps.h>
 
@@ -15,10 +16,10 @@ inline void fill_inplace(Tensor& self, const Scalar& value_scalar) {
 
 namespace detail {
 Tensor& scalar_fill(Tensor& self, const Scalar& value) {
-  AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND4(
-      kComplexHalf, kHalf, kBool, kBFloat16, self.scalar_type(), "fill_out", [&]() {
+  AT_DISPATCH_V2(
+      self.scalar_type(), "fill_out", AT_WRAP([&]() {
         fill_inplace<scalar_t>(self, value);
-      });
+      }), kComplexHalf, kHalf, kBool, kBFloat16, AT_EXPAND(AT_ALL_TYPES_AND_COMPLEX), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES));
   return self;
 }
 

aten/src/ATen/ScalarOps.h

@@ -33,27 +33,9 @@ inline at::Tensor scalar_to_tensor(
     const Device device = at::kCPU) {
   // This is the fast track we have for CPU scalar tensors.
   if (device == at::kCPU) {
-    if (s.isFloatingPoint()) {
-      return at::detail::scalar_tensor_static(s, at::kDouble, at::kCPU);
-    } else if (s.isComplex()) {
-      return at::detail::scalar_tensor_static(s, at::kComplexDouble, at::kCPU);
-    } else if (s.isBoolean()) {
-      return at::detail::scalar_tensor_static(s, at::kBool, at::kCPU);
-    } else {
-      AT_ASSERT(s.isIntegral(false));
-      return at::detail::scalar_tensor_static(s, at::kLong, at::kCPU);
-    }
-  }
-  if (s.isFloatingPoint()) {
-    return at::scalar_tensor(s, at::device(device).dtype(at::kDouble));
-  } else if (s.isBoolean()) {
-    return at::scalar_tensor(s, at::device(device).dtype(at::kBool));
-  } else if (s.isComplex()) {
-    return at::scalar_tensor(s, at::device(device).dtype(at::kComplexDouble));
-  } else {
-    AT_ASSERT(s.isIntegral(false));
-    return at::scalar_tensor(s, at::device(device).dtype(at::kLong));
+    return at::detail::scalar_tensor_static(s, s.type(), at::kCPU);
   }
+  return at::scalar_tensor(s, at::device(device).dtype(s.type()));
 }
 
 } // namespace c10

aten/src/ATen/native/TensorFactories.h

@@ -4,6 +4,7 @@
 #include <ATen/EmptyTensor.h>
 #include <ATen/TensorIterator.h>
 #include <ATen/Dispatch.h>
+#include <ATen/Dispatch_v2.h>
 #include <ATen/native/DispatchStub.h>
 
 #ifndef AT_PER_OPERATOR_HEADERS
@@ -107,10 +108,10 @@ inline Tensor& fill_empty_deterministic_(Tensor& tensor) {
         tensor.fill_(std::numeric_limits<scalar_t>::quiet_NaN());
     });
   } else {
-    AT_DISPATCH_INTEGRAL_TYPES_AND(
-      kBool, tensor.scalar_type(), "fill_empty_deterministic_", [&]() {
+    AT_DISPATCH_V2(
+      tensor.scalar_type(), "fill_empty_deterministic_", AT_WRAP([&]() {
         tensor.fill_(std::numeric_limits<scalar_t>::max());
-    });
+    }), kBool, AT_EXPAND(AT_INTEGRAL_TYPES_V2));
   }
   return tensor;
 }

test/test_torch.py

@@ -1268,7 +1268,7 @@ def test_nondeterministic_resize_quantized(self, device, dtype):
 
     @skipXLA
     @skipIfTorchInductor("https://github.com/pytorch/pytorch/issues/113707")
-    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
+    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16, torch.uint16, torch.uint32, torch.uint64))
     def test_deterministic_resize(self, device, dtype):
         test_cases = [
             # size, stride, resize_size

torch/csrc/utils/pybind.cpp

@@ -127,7 +127,11 @@ py::handle type_caster<c10::Scalar>::cast(
     if (scalar.isSymbolic()) {
       return py::cast(scalar.toSymInt()).release();
     } else {
-      return py::cast(scalar.toLong()).release();
+      if (scalar.type() == at::ScalarType::UInt64) {
+        return py::cast(scalar.toUInt64()).release();
+      } else {
+        return py::cast(scalar.toLong()).release();
+      }
     }
   } else if (scalar.isFloatingPoint()) {
     // This isn't strictly necessary but we add it for symmetry

torch/csrc/utils/python_arg_parser.cpp

@@ -1649,7 +1649,21 @@ at::Tensor PythonArgs::tensor_slow(int i) {
   if (PyBool_Check(obj)) {
     scalar = at::Scalar(THPUtils_unpackBool(obj));
   } else if (THPUtils_checkLong(obj)) {
-    scalar = at::Scalar(THPUtils_unpackLong(obj));
+    int overflow = -1;
+    long long value = PyLong_AsLongLongAndOverflow(obj, &overflow);
+    if (value == -1 && PyErr_Occurred()) {
+      throw python_error();
+    }
+    if (overflow != 0) {
+      // try unsigned
+      unsigned long long value = PyLong_AsUnsignedLongLong(obj);
+      if (value == static_cast<unsigned long long>(-1) && PyErr_Occurred()) {
+        throw python_error();
+      }
+      scalar = at::Scalar(static_cast<uint64_t>(value));
+    } else {
+      scalar = at::Scalar(static_cast<int64_t>(value));
+    }
   } else if (PyComplex_Check(obj)) {
     scalar = at::Scalar(THPUtils_unpackComplexDouble(obj));
   } else if (THPUtils_checkDouble(obj)) {
@@ -1712,7 +1726,21 @@ at::Scalar PythonArgs::scalar_slow(PyObject* arg) {
   }
 
   if (THPUtils_checkLong(arg)) {
-    return at::Scalar(static_cast<int64_t>(THPUtils_unpackLong(arg)));
+    int overflow = -1;
+    long long value = PyLong_AsLongLongAndOverflow(arg, &overflow);
+    if (value == -1 && PyErr_Occurred()) {
+      throw python_error();
+    }
+    if (overflow != 0) {
+      // try unsigned
+      unsigned long long value = PyLong_AsUnsignedLongLong(arg);
+      if (value == static_cast<unsigned long long>(-1) && PyErr_Occurred()) {
+        throw python_error();
+      }
+      return at::Scalar(static_cast<uint64_t>(value));
+    } else {
+      return at::Scalar(static_cast<int64_t>(value));
+    }
   }
 
   if (PyBool_Check(arg)) {