support multiple template parameter in KernelType for REGISTER_OP_XPU…

…_KERNEL (#2932)
PaddlePaddle · Jul 18, 2017 · 051676a · 051676a
1 parent 861b66d
commit 051676a
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git a/paddle/framework/op_registry.h b/paddle/framework/op_registry.h
@@ -311,7 +311,7 @@ class OpRegisterHelper {
 /**
  * Macro to Register OperatorKernel.
  */
-#define REGISTER_OP_KERNEL(type, DEVICE_TYPE, PlaceType, KernelType)      \
+#define REGISTER_OP_KERNEL(type, DEVICE_TYPE, PlaceType, ...)             \
   STATIC_ASSERT_GLOBAL_NAMESPACE(                                         \
       __reg_op_kernel_##type##_##DEVICE_TYPE##__,                         \
       "REGISTER_OP_KERNEL must be in global namespace");                  \
@@ -320,17 +320,19 @@ class OpRegisterHelper {
       ::paddle::framework::OperatorWithKernel::OpKernelKey key;           \
       key.place_ = PlaceType();                                           \
       ::paddle::framework::OperatorWithKernel::AllOpKernels()[#type][key] \
-          .reset(new KernelType());                                       \
+          .reset(new __VA_ARGS__());                                      \
     }                                                                     \
   };                                                                      \
   static __op_kernel_register__##type##__ __reg_kernel_##type##__;        \
   int __op_kernel_register_##type##_handle_##DEVICE_TYPE##__() { return 0; }
 
-#define REGISTER_OP_GPU_KERNEL(type, KernelType) \
-  REGISTER_OP_KERNEL(type, GPU, ::paddle::platform::GPUPlace, KernelType)
+// (type, KernelType)
+#define REGISTER_OP_GPU_KERNEL(type, ...) \
+  REGISTER_OP_KERNEL(type, GPU, ::paddle::platform::GPUPlace, __VA_ARGS__)
 
-#define REGISTER_OP_CPU_KERNEL(type, KernelType) \
-  REGISTER_OP_KERNEL(type, CPU, ::paddle::platform::CPUPlace, KernelType)
+// (type, KernelType)
+#define REGISTER_OP_CPU_KERNEL(type, ...) \
+  REGISTER_OP_KERNEL(type, CPU, ::paddle::platform::CPUPlace, __VA_ARGS__)
 
 /**
  * Macro to mark what Operator and Kernel we will use and tell the compiler to

diff --git a/paddle/framework/operator_test.cc b/paddle/framework/operator_test.cc
@@ -102,6 +102,7 @@ class OpWithKernelTest : public OperatorWithKernel {
                   const std::vector<Tensor*>& outputs) const override {}
 };
 
+template <typename T1, typename T2>
 class CPUKernelTest : public OpKernel {
  public:
   void Compute(const KernelContext& ctx) const {
@@ -171,7 +172,8 @@ class CPUKernalMultiInputsTest : public OpKernel {
 
 REGISTER_OP(op_with_kernel, paddle::framework::OpWithKernelTest,
             paddle::framework::OpKernelTestProtoAndCheckerMaker);
-REGISTER_OP_CPU_KERNEL(op_with_kernel, paddle::framework::CPUKernelTest);
+REGISTER_OP_CPU_KERNEL(op_with_kernel,
+                       paddle::framework::CPUKernelTest<float, float>);
 
 // test with single input
 TEST(OpKernel, all) {