Improve compare kernel

aten/src/ATen/native/TensorIterator.cpp

@@ -230,15 +230,18 @@ void TensorIterator::compute_types() {
   }
 
   for (auto &op : operands_) {
+    bool skip_output = compute_common_dtype_only_for_inputs && op.is_output;
+    bool is_different = op.tensor.defined() && op.current_dtype != common_dtype_;
+
     if (may_have_differing_types) {
       validate_dtype(op, common_dtype_, common_dtype_strategy_);
-      bool cast_by_copy = compute_common_dtype && !common_device_is_cuda && (!compute_common_dtype_only_for_inputs || !op.is_output);
+      bool cast_by_copy = compute_common_dtype && !common_device_is_cuda && !skip_output;
       if (cast_by_copy) {
         maybe_copy_casting_to_common_dtype(op, common_dtype_);
       }
     }
 
-    if (op.tensor.defined() && op.current_dtype != common_dtype_) {
+    if (is_different && !skip_output) {
       have_differing_types_ = true;
     }
 
@@ -686,6 +689,7 @@ TensorIterator TensorIterator::comparison_op(Tensor& out, const Tensor& a,
   iter.allow_cpu_scalars_ = true;
   iter.compute_common_dtype_only_for_inputs();
   iter.build();
+  iter.dynamic_cast_if(iter.dtype() != kBool);
   return iter;
 }
 

aten/src/ATen/native/TensorIterator.h

@@ -301,7 +301,7 @@ struct CAFFE2_API TensorIterator {
   bool is_final_output() const { return final_output_; }
 
   bool needs_dynamic_casting() const {
-    return (common_dtype_strategy_ != CommonDTypeStrategy::NONE) && have_differing_types_;
+    return force_dynamic_casting_ || ((common_dtype_strategy_ != CommonDTypeStrategy::NONE) && have_differing_types_);
   }
 
   void set_check_mem_overlap(bool check_mem_overlap) {
@@ -343,6 +343,10 @@ struct CAFFE2_API TensorIterator {
     resize_outputs_ = false;
   }
 
+  void dynamic_cast_if(bool condition) {
+    force_dynamic_casting_ = force_dynamic_casting_ || condition;
+  }
+
   void build();
 
 protected:
@@ -383,6 +387,7 @@ struct CAFFE2_API TensorIterator {
   bool final_output_ = true;
   bool check_mem_overlap_ = false;
   bool have_differing_types_ = false;
+  bool force_dynamic_casting_ = false;
   bool all_ops_same_shape_ = false;
   bool requires_channels_last_output_ = false;
 };

aten/src/ATen/native/cuda/BinaryCompareKernel.cu

@@ -11,47 +11,47 @@ namespace at { namespace native {
 
 void lt_kernel_cuda(TensorIterator& iter) {
   AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.common_dtype(), "lt_cuda", [&]() {
-    gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t {
+    gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool {
       return a < b;
     });
   });
 }
 
 void le_kernel_cuda(TensorIterator& iter) {
   AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.common_dtype(), "le_cuda", [&]() {
-    gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t {
+    gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool {
       return a <= b;
     });
   });
 }
 
 void gt_kernel_cuda(TensorIterator& iter) {
   AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.common_dtype(), "gt_cuda", [&]() {
-    gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t {
+    gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool {
       return a > b;
     });
   });
 }
 
 void ge_kernel_cuda(TensorIterator& iter) {
   AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.common_dtype(), "ge_cuda", [&]() {
-    gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t {
+    gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool {
       return a >= b;
     });
   });
 }
 
 void eq_kernel_cuda(TensorIterator& iter) {
   AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.common_dtype(), "eq_cuda", [&]() {
-    gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t {
+    gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool {
       return a == b;
     });
   });
 }
 
 void ne_kernel_cuda(TensorIterator& iter) {
   AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.common_dtype(), "ne_cuda", [&]() {
-    gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t {
+    gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool {
       return a != b;
     });
   });

aten/src/ATen/native/cuda/BinaryMiscOpsKernels.cu

@@ -40,8 +40,8 @@ void bitwise_xor_kernel_cuda(TensorIterator& iter) {
 
 void logical_xor_kernel_cuda(TensorIterator& iter) {
   AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBool, iter.common_dtype(), "logical_xor_cuda", [&]() {
-    gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t {
-      return static_cast<scalar_t>(bool(a) != bool(b));
+    gpu_kernel_with_scalars(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> bool {
+      return bool(a) != bool(b);
     });
   });
 }

aten/src/ATen/native/cuda/Copy.cu

@@ -58,9 +58,7 @@ 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.promote_common_dtype();
+    iter.dynamic_cast_if(true);
     AT_DISPATCH_ALL_TYPES_AND3(kHalf, kBool, kBFloat16, iter.dtype(0), "copy_", [&] {
       gpu_kernel(iter, []GPU_LAMBDA(scalar_t x) { return x; });
     });