[pten] update isnan registration (#39419)

* update isnan registration * fix compile
PaddlePaddle · Feb 10, 2022 · 14ed2f5 · 14ed2f5
1 parent c7c1db3
commit 14ed2f5
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 17 deletions.
diff --git a/paddle/fluid/operators/isfinite_op.cc b/paddle/fluid/operators/isfinite_op.cc
@@ -132,4 +132,33 @@ namespace ops = paddle::operators;
 REGISTER_OP_MAKER(isinf, "isinf(X)");
 REGISTER_OP_MAKER(isnan, "isnan(X)");
 REGISTER_OP_MAKER(isfinite, "isfinite(X)");
-FOR_EACH_KERNEL_FUNCTOR(REGISTER_OVERFLOW_CPU_KERNEL);
+
+REGISTER_OP_CPU_KERNEL(isinf,
+                       ops::OverflowKernel<paddle::platform::CPUDeviceContext,
+                                           int, ops::InfinityFunctor>,
+                       ops::OverflowKernel<paddle::platform::CPUDeviceContext,
+                                           int64_t, ops::InfinityFunctor>,
+                       ops::OverflowKernel<paddle::platform::CPUDeviceContext,
+                                           float, ops::InfinityFunctor>,
+                       ops::OverflowKernel<paddle::platform::CPUDeviceContext,
+                                           double, ops::InfinityFunctor>);
+
+REGISTER_OP_CPU_KERNEL(isnan,
+                       ops::OverflowKernel<paddle::platform::CPUDeviceContext,
+                                           int, ops::NANFunctor>,
+                       ops::OverflowKernel<paddle::platform::CPUDeviceContext,
+                                           int64_t, ops::NANFunctor>,
+                       ops::OverflowKernel<paddle::platform::CPUDeviceContext,
+                                           float, ops::NANFunctor>,
+                       ops::OverflowKernel<paddle::platform::CPUDeviceContext,
+                                           double, ops::NANFunctor>);
+
+REGISTER_OP_CPU_KERNEL(isfinite,
+                       ops::OverflowKernel<paddle::platform::CPUDeviceContext,
+                                           int, ops::IsfiniteFunctor>,
+                       ops::OverflowKernel<paddle::platform::CPUDeviceContext,
+                                           int64_t, ops::IsfiniteFunctor>,
+                       ops::OverflowKernel<paddle::platform::CPUDeviceContext,
+                                           float, ops::IsfiniteFunctor>,
+                       ops::OverflowKernel<paddle::platform::CPUDeviceContext,
+                                           double, ops::IsfiniteFunctor>);
diff --git a/paddle/fluid/operators/isfinite_op.cu b/paddle/fluid/operators/isfinite_op.cu
@@ -17,15 +17,32 @@
 namespace ops = paddle::operators;
 namespace plat = paddle::platform;
 
-#define REGISTER_OVERFLOW_CUDA_KERNEL(op_type, functor)                       \
-  REGISTER_OP_CUDA_KERNEL(                                                    \
-      op_type, ops::OverflowKernel<paddle::platform::CUDADeviceContext, int,  \
-                                   ops::functor>,                             \
-      ops::OverflowKernel<paddle::platform::CUDADeviceContext, float,         \
-                          ops::functor>,                                      \
-      ops::OverflowKernel<paddle::platform::CUDADeviceContext, double,        \
-                          ops::functor>,                                      \
-      ops::OverflowKernel<paddle::platform::CUDADeviceContext, plat::float16, \
-                          ops::functor>);
+REGISTER_OP_CUDA_KERNEL(
+    isinf, ops::OverflowKernel<paddle::platform::CUDADeviceContext, int,
+                               ops::InfinityFunctor>,
+    ops::OverflowKernel<paddle::platform::CUDADeviceContext, float,
+                        ops::InfinityFunctor>,
+    ops::OverflowKernel<paddle::platform::CUDADeviceContext, double,
+                        ops::InfinityFunctor>,
+    ops::OverflowKernel<paddle::platform::CUDADeviceContext, plat::float16,
+                        ops::InfinityFunctor>);
 
-FOR_EACH_KERNEL_FUNCTOR(REGISTER_OVERFLOW_CUDA_KERNEL);
+REGISTER_OP_CUDA_KERNEL(isnan,
+                        ops::OverflowKernel<paddle::platform::CUDADeviceContext,
+                                            int, ops::NANFunctor>,
+                        ops::OverflowKernel<paddle::platform::CUDADeviceContext,
+                                            float, ops::NANFunctor>,
+                        ops::OverflowKernel<paddle::platform::CUDADeviceContext,
+                                            double, ops::NANFunctor>,
+                        ops::OverflowKernel<paddle::platform::CUDADeviceContext,
+                                            plat::float16, ops::NANFunctor>);
+
+REGISTER_OP_CUDA_KERNEL(
+    isfinite, ops::OverflowKernel<paddle::platform::CUDADeviceContext, int,
+                                  ops::IsfiniteFunctor>,
+    ops::OverflowKernel<paddle::platform::CUDADeviceContext, float,
+                        ops::IsfiniteFunctor>,
+    ops::OverflowKernel<paddle::platform::CUDADeviceContext, double,
+                        ops::IsfiniteFunctor>,
+    ops::OverflowKernel<paddle::platform::CUDADeviceContext, plat::float16,
+                        ops::IsfiniteFunctor>);
diff --git a/paddle/fluid/operators/isfinite_op.h b/paddle/fluid/operators/isfinite_op.h
@@ -73,8 +73,3 @@ class OverflowKernel : public framework::OpKernel<T> {
 
 }  // namespace operators
 }  // namespace paddle
-
-#define FOR_EACH_KERNEL_FUNCTOR(__macro) \
-  __macro(isinf, InfinityFunctor);       \
-  __macro(isnan, NANFunctor);            \
-  __macro(isfinite, IsfiniteFunctor);