Global function, op_support_gpu (#4980)

paddle/framework/operator.cc

@@ -252,5 +252,20 @@ std::ostream& operator<<(std::ostream& os,
   return os;
 }
 
+bool OpSupportGPU(const std::string& op_type) {
+  auto& all_kernels = OperatorWithKernel::AllOpKernels();
+  auto it = all_kernels.find(op_type);
+  if (it == all_kernels.end()) {
+    // All control operator must support GPU
+    return true;
+  }
+  for (auto& kern_pair : it->second) {
+    if (platform::is_gpu_place(kern_pair.first.place_)) {
+      return true;
+    }
+  }
+  return false;
+}
+
 }  // namespace framework
 }  // namespace paddle

paddle/framework/operator.h

@@ -649,5 +649,7 @@ class OperatorWithKernel : public OperatorBase {
 std::ostream& operator<<(std::ostream& os,
                          const OperatorWithKernel::OpKernelKey& kernel_key);
 
+extern bool OpSupportGPU(const std::string& op_type);
+
 }  // namespace framework
 }  // namespace paddle

paddle/pybind/pybind.cc

@@ -466,6 +466,8 @@ All parameter, weight, gradient are variables in Paddle.
   BindVarDsec(m);
   BindOpDesc(m);
 
+  m.def("op_support_gpu", OpSupportGPU);
+
   return m.ptr();
 }
 }  // namespace pybind

python/paddle/v2/framework/tests/test_op_support_gpu.py

@@ -0,0 +1,11 @@
+import unittest
+import paddle.v2.framework.core as core
+
+
+class TestOpSupportGPU(unittest.TestCase):
+    def test_case(self):
+        self.assertEqual(core.is_compile_gpu(), core.op_support_gpu("sum"))
+
+
+if __name__ == '__main__':
+    unittest.main()