Update on "Simplify copy kernel"

zasdfgbnm · zasdfgbnm · commit dd612d04a33c · 2019-10-25T15:43:18.000-07:00
Using the new type promotion and dynamic casting added to `TensorIterator`, the copy kernels could be greatly simplified. For benchmark, see #28352 (comment) [ghstack-poisoned]
diff --git a/aten/src/ATen/native/TensorIterator.cpp b/aten/src/ATen/native/TensorIterator.cpp
@@ -189,7 +189,7 @@ void TensorIterator::compute_types() {
     auto operands = compute_common_dtype_only_for_inputs ? at::ArrayRef<OperandInfo>(operands_).slice(noutputs()) : operands_;
     auto common_type = compute_common_type_(operands);
     auto common_device = std::get<0>(common_type);
-    bool common_device_is_cuda = common_device.is_cuda();
+    common_device_is_cuda = common_device.is_cuda();
     common_dtype_ = std::get<1>(common_type);
     may_have_differing_types_ = !std::get<2>(common_type);
     bool has_cpu_scalar = false;
diff --git a/aten/src/ATen/native/cuda/Copy.cu b/aten/src/ATen/native/cuda/Copy.cu
@@ -58,7 +58,10 @@ static void copy_device_to_device(TensorIterator& iter, bool non_blocking) {
         cudaMemcpyDeviceToDevice,
         copy_stream));
   } else {
+    // this is done intentionally done after build because copy has a "promotion"
+    // rule that always "promote" to target dtype.
     iter.set_common_dtype(iter.dtype());
+    iter.promote_common_dtype();
     AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.dtype(0), "copy_", [&] {
       gpu_kernel(iter, []GPU_LAMBDA(scalar_t x) { return x; });
     });
