Avoid double dispatch in logical_not for compilation speed reasons. (#38565)

gchanan · facebook-github-bot · commit b27be3e0c59e · 2020-05-18T09:25:54.000-07:00
Summary: Pull Request resolved: #38565 Also note this turns on "-Wno-unused-local-typedefs" because we are using dispatch macros for error checking. Test Plan: Imported from OSS Differential Revision: D21598478 Pulled By: gchanan fbshipit-source-id: 28f9ad01bd678df0601a10d0daf3ed31c47c4ab2
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -504,6 +504,7 @@ if(NOT MSVC)
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-variable")
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-function")
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-result")
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-unused-local-typedefs")
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-strict-overflow")
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-strict-aliasing")
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-error=deprecated-declarations")
diff --git a/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp b/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp
@@ -129,6 +129,9 @@ static void frac_kernel(TensorIterator& iter) {
 }
 
 static void logical_not_kernel(TensorIterator& iter) {
+  // NOTE: this implementation differs from the CUDA implementation which only does single dispatch
+  // (to avoid expensive compilation) because CPU kernels don't handle dynamic_casting
+  // (see needs_dynamic_casting).
   AT_DISPATCH_ALL_TYPES_AND2(kBool, kHalf, iter.dtype(1), "logical_not_cpu", [&]() {
     using self_t = scalar_t;
     AT_DISPATCH_ALL_TYPES_AND2(kBool, kHalf, iter.dtype(0), "logical_not_cpu", [&]() {
diff --git a/aten/src/ATen/native/cuda/UnarySignKernels.cu b/aten/src/ATen/native/cuda/UnarySignKernels.cu
@@ -11,11 +11,12 @@
 namespace at { namespace native {
 
 void logical_not_kernel_cuda(TensorIterator& iter) {
+  // error check -- this is just ensuring we don't dispatch on types that aren't in ALL_TYPES_AND2(...)
+  // so we don't have to maintain a separate list or to do double dispatch.
+  AT_DISPATCH_ALL_TYPES_AND2(kBool, kHalf, iter.dtype(0), "logical_not_cuda", [&]() {});
+
   AT_DISPATCH_ALL_TYPES_AND2(kBool, kHalf, iter.dtype(1), "logical_not_cuda", [&]() {
-    using self_t = scalar_t;
-    AT_DISPATCH_ALL_TYPES_AND2(kBool, kHalf, iter.dtype(0), "logical_not_cuda", [&]() {
-      gpu_kernel(iter, []GPU_LAMBDA(self_t a) -> scalar_t { return static_cast<scalar_t>(!a); });
-    });
+    gpu_kernel(iter, []GPU_LAMBDA(scalar_t a) -> bool { return !a; });
   });
 }
 

Original file line number	Diff line number	Diff line change
`@@ -129,6 +129,9 @@ static void frac_kernel(TensorIterator& iter) {`
`129`	`129`	`}`
`130`	`130`
`131`	`131`	`static void logical_not_kernel(TensorIterator& iter) {`
	`132`	`+ // NOTE: this implementation differs from the CUDA implementation which only does single dispatch`
	`133`	`+ // (to avoid expensive compilation) because CPU kernels don't handle dynamic_casting`
	`134`	`+ // (see needs_dynamic_casting).`
`132`	`135`	`AT_DISPATCH_ALL_TYPES_AND2(kBool, kHalf, iter.dtype(1), "logical_not_cpu", [&]() {`
`133`	`136`	`using self_t = scalar_t;`
`134`	`137`	`AT_DISPATCH_ALL_TYPES_AND2(kBool, kHalf, iter.dtype(0), "logical_not_cpu", [&]() {`