Changes TensorIterator computation to not consider out kwarg, lets UnaryOps safe cast to out (pytorch#39655)

Summary:
**BC breaking note:**

In PyTorch 1.5 passing the out= kwarg to some functions, like torch.add, could affect the computation. That is,

```
out = torch.add(a, b)
```

could produce a different tensor than

```
torch.add(a, b, out=out)
```

This is because previously the out argument participated in the type promotion rules. For greater consistency with NumPy, Python, and C++, in PyTorch 1.6 the out argument no longer participates in type promotion, and has no effect on the computation performed.

**ORIGINAL PR NOTE**

This PR effectively rewrites Tensor Iterator's "compute_types" function to both clarify its behavior and change how our type promotion works to never consider the out argument when determining the iterator's "common dtype," AKA its "computation type." That is,

```
a = op(b, c)
```

should always produce the same result as

```
op(b, c, out=a)
```

This is consistent with NumPy and programming languages like Python and C++.

The conceptual model for this change is that a TensorIterator may have a "common computation type" that all inputs are cast to and its computation performed in. This common computation type, if it exists, is determined by applying our type promotion rules to the inputs.

A common computation type is natural for some classes of functions, like many binary elementwise functions (e.g. add, sub, mul, div...). (NumPy describes these as "universal functions.") Many functions, however, like indexing operations, don't have a natural common computation type. In the future we'll likely want to support setting the TensorIterator's common computation type explicitly to enable "floating ufuncs" like the sin function that promote integer types to the default scalar type. Logic like that is beyond the type promotion system, which can only review inputs.

Implementing this change in a readable and maintainable manner was challenging because compute_types() has had many small modifications from many authors over ~2 year period, and the existing logic was in some places outdated and in other places unnecessarily complicated. The existing "strategies" approach also painted with a broad brush, and two of them no longer made conceptual sense after this change. As a result, the new version of this function has a small set of flags to control its behavior. This has the positive effect of disentangling checks like all operands having the same device and their having the same dtype.

Additional changes in this PR:

- Unary operations now support out arguments with different dtypes. Like binary ops they check canCast(computation type, out dtype).
- The dtype checking for lerp was outdated and its error message included the wrong variable. It has been fixed.
- The check for whether all tensors are on the same device has been separated from other checks. TensorIterators used by copy disable this check.
- As a result of this change, the output dtype can be computed if only the input types are available.
- The "fast path" for checking if a common dtype computation is necessary has been updated and simplified to also handle zero-dim tensors.
- A couple helper functions for compute_types() have been inlined to improve readability.
- The confusingly named and no longer used promote_gpu_output_dtypes_ has been removed. This variable was intended to support casting fp16 reductions on GPU, but it has become a nullop. That logic is now implemented here: https://github.com/pytorch/pytorch/blob/856215509d89c935cd1768ce4b496d4fc0e919a6/aten/src/ATen/native/ReduceOpsUtils.h#L207.
Pull Request resolved: pytorch#39655

Differential Revision: D21970878

Pulled By: mruberry

fbshipit-source-id: 5e6354c78240877ab5d6b1f7cfb351bd89049012

Author: Mike Ruberry

aten/src/ATen/native/BinaryOps.cpp

@@ -111,23 +111,34 @@ Tensor& remainder_(Tensor& self, const Tensor& other) {
 }
 
 Tensor& true_divide_out(Tensor& result, const Tensor& self, const Tensor& divisor) {
-  TORCH_CHECK(!isIntegralType(result.scalar_type(), /*includeBool=*/ true),
-            "True division requires a floating output type, but got ",
-            result.scalar_type());
+  // If both inputs have integral (or bool) types, creates
+  // temporary float copies as new inputs.
+  if (isIntegralType(self.scalar_type(), /*includeBool=*/ true)
+   && isIntegralType(divisor.scalar_type(), /*includeBool=*/ true)) {
+    const auto scalar_type = typeMetaToScalarType(c10::get_default_dtype());
+    auto iter = TensorIterator::binary_op(result,
+                                          self.to(scalar_type),
+                                          divisor.to(scalar_type),
+                                          /*check_mem_overlap=*/ true);
+    div_stub(iter.device_type(), iter);
+    return result;
+  }
   auto iter = TensorIterator::binary_op(result, self, divisor, /*check_mem_overlap=*/ true);
   div_stub(iter.device_type(), iter);
   return result;
 }
 
 Tensor true_divide(const Tensor& self, const Tensor& divisor) {
-  // If both inputs have integral (or bool) types, sets the output to have
-  // the default (floating) scalar type
+  // If both inputs have integral (or bool) types, creates
+  // temporary float copies as new inputs and sets the result's type to
+  // the default scalar type
   if (isIntegralType(self.scalar_type(), /*includeBool=*/ true)
    && isIntegralType(divisor.scalar_type(), /*includeBool=*/ true)) {
     const auto scalar_type = typeMetaToScalarType(c10::get_default_dtype());
     Tensor result = at::empty({0}, self.options().dtype(scalar_type));
-
-    auto iter = TensorIterator::binary_op(result, self, divisor);
+    auto iter = TensorIterator::binary_op(result,
+                                          self.to(scalar_type),
+                                          divisor.to(scalar_type));
     div_stub(iter.device_type(), iter);
     return result;
   }

aten/src/ATen/native/Copy.cpp

@@ -140,7 +140,8 @@ static Tensor & copy_impl(Tensor & self, const Tensor & src, bool non_blocking)
   iter.add_output(self);
   iter.add_input(src);
   iter.dont_resize_outputs();
-  iter.dont_compute_common_dtype();
+  iter.check_all_same_dtype(false);
+  iter.check_all_same_device(false);
   iter.build();
 
   if (iter.numel() == 0) {

aten/src/ATen/native/ReduceOpsUtils.h

@@ -212,7 +212,7 @@ static TensorIterator make_reduction(
 
 static TensorIterator make_reduction(
     const char* name, Tensor& result1, Tensor& result2, const Tensor& self, IntArrayRef dim,
-    bool keepdim, ScalarType dtype1, ScalarType dtype2, bool promote_gpu_output_dtypes=true)
+    bool keepdim, ScalarType dtype1, ScalarType dtype2)
 {
   // check that result type and dtype match if provided
   TORCH_CHECK(
@@ -240,9 +240,9 @@ static TensorIterator make_reduction(
   // product of templated kernel launches.
   if (self.scalar_type() == dtype1 ||
       (self.is_cuda() && self.scalar_type() == kHalf && dtype1 == kFloat)) {
-    return TensorIterator::reduce_op(viewed_result1, viewed_result2, self, promote_gpu_output_dtypes);
+    return TensorIterator::reduce_op(viewed_result1, viewed_result2, self);
   }
-  return TensorIterator::reduce_op(viewed_result1, viewed_result2, self.to(dtype1), promote_gpu_output_dtypes);
+  return TensorIterator::reduce_op(viewed_result1, viewed_result2, self.to(dtype1));
 }
 
 static TensorIterator make_reduction(

aten/src/ATen/native/TensorAdvancedIndexing.cpp

@@ -215,11 +215,15 @@ static AdvancedIndex make_info(Tensor self, TensorList orig) {
 static TensorIterator make_index_put_iterator(const AdvancedIndex& info, const Tensor& value) {
   TORCH_CHECK(is_expandable_to(value.sizes(), info.src.sizes()), "shape mismatch: value tensor of shape ", value.sizes(),
              " cannot be broadcast to indexing result of shape ", info.src.sizes());
+  TORCH_CHECK(value.scalar_type() == info.src.scalar_type(),
+              "Index put requires the source and destination dtypes match, "
+              "got ", info.src.scalar_type(), " for the destination "
+              "and ", value.scalar_type(), " for the source.");
   auto iter = TensorIterator();
-  iter.dont_compute_common_dtype();
   iter.dont_resize_outputs();
+  iter.check_all_same_dtype(false);
   iter.add_output(info.src);
-  iter.add_input(value, info.src.device(), info.src.scalar_type());
+  iter.add_input(value);
   for (auto& index : info.indices) {
     iter.add_input(index);
   }
@@ -229,7 +233,7 @@ static TensorIterator make_index_put_iterator(const AdvancedIndex& info, const T
 
 static TensorIterator make_index_iterator(const AdvancedIndex& info) {
   auto iter = TensorIterator();
-  iter.dont_compute_common_dtype();
+  iter.check_all_same_dtype(false);
   iter.add_output(Tensor(), info.src.device(), info.src.scalar_type());
   iter.add_input(info.src);
   for (auto& index : info.indices) {
@@ -241,7 +245,7 @@ static TensorIterator make_index_iterator(const AdvancedIndex& info) {
 
 static TensorIterator make_index_out_iterator(const AdvancedIndex& info, Tensor& result) {
   auto iter = TensorIterator();
-  iter.dont_compute_common_dtype();
+  iter.check_all_same_dtype(false);
   iter.add_output(result, info.src.device(), info.src.scalar_type());
   iter.add_input(info.src);
   for (auto& index : info.indices) {
@@ -437,7 +441,7 @@ Tensor & index_select_out_cpu_(Tensor & result, const Tensor & self, int64_t dim
     auto slice_size = selfSlice.numel();
 
     auto iter = TensorIterator();
-    iter.dont_compute_common_dtype();
+    iter.check_all_same_dtype(false);
     iter.dont_resize_outputs();
     iter.add_output(resultSlice);
     iter.add_input(selfSlice);
@@ -571,7 +575,7 @@ static Tensor & masked_fill_impl_cpu(Tensor & self, const Tensor & mask, Scalar
   }
 
   auto iter = TensorIterator();
-  iter.dont_compute_common_dtype();
+  iter.check_all_same_dtype(false);
   iter.dont_resize_outputs();
   iter.add_output(self);
   iter.add_input(mask);
@@ -658,7 +662,7 @@ static Tensor & masked_select_out_impl_cpu(Tensor & result, const Tensor & self,
   bool use_serial_kernel = self.numel() < at::internal::GRAIN_SIZE || at::get_num_threads() == 1;
   if (use_serial_kernel) {
     auto iter = TensorIterator();
-    iter.dont_compute_common_dtype();
+    iter.check_all_same_dtype(false);
     iter.dont_resize_outputs();
     iter.add_output(result_strided);
     iter.add_input(_self);
@@ -681,7 +685,7 @@ static Tensor & masked_select_out_impl_cpu(Tensor & result, const Tensor & self,
   std::partial_sum(mask_long_data, mask_long_data + mask_long.numel(), mask_prefix_sum_data);
 
   auto iter = TensorIterator();
-  iter.dont_compute_common_dtype();
+  iter.check_all_same_dtype(false);
   iter.dont_resize_outputs();
   iter.add_output(result_strided);
   iter.add_input(_self);

aten/src/ATen/native/TensorCompare.cpp

@@ -151,12 +151,12 @@ Tensor _s_where(const Tensor& condition, const Tensor& self, const Tensor& other
   TORCH_CHECK(self.dtype() == other.dtype(), "expected scalar type ", self.dtype(), " but found ", other.dtype());
   Tensor ret = at::empty(self.sizes(), self.options());
   auto iter = at::TensorIterator();
+  iter.check_all_same_dtype(false);
   iter.set_check_mem_overlap(true);
   iter.add_output(ret);
   iter.add_input(condition);
   iter.add_input(self);
   iter.add_input(other);
-  iter.dont_compute_common_dtype();
   iter.build();
   where_kernel(iter.device_type(), iter, condition.scalar_type());
   return ret;