From 493396d81cbcd1e29b5ea6c3aa11cfa20496b773 Mon Sep 17 00:00:00 2001
From: Qiao Longfei <qiaolongfei@baidu.com>
Date: Mon, 7 Aug 2017 19:50:37 +0800
Subject: [PATCH] add support_gpu (#3304)

* add support_gpu

* fix allclose

* fix name error and symplify code
---
 paddle/framework/op_registry.h                   | 6 ++++++
 paddle/framework/pybind.cc                       | 2 ++
 python/paddle/v2/framework/tests/op_test_util.py | 6 ++++--
 3 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/paddle/framework/op_registry.h b/paddle/framework/op_registry.h
index 6c26183818a..b2813da83d9 100644
--- a/paddle/framework/op_registry.h
+++ b/paddle/framework/op_registry.h
@@ -260,6 +260,12 @@ class OpRegistry {
     return CreateOp(op_desc.type(), inputs, outputs, attrs);
   }
 
+  static bool SupportGPU(const std::string& op_type) {
+    OperatorWithKernel::OpKernelKey key;
+    key.place_ = platform::GPUPlace();
+    return OperatorWithKernel::AllOpKernels().at(op_type).count(key) != 0;
+  }
+
   static std::shared_ptr<OperatorBase> CreateGradOp(const OperatorBase& op) {
     PADDLE_ENFORCE(!op.IsNetOp(),
                    "Use framework::Backward to get backward ops");
diff --git a/paddle/framework/pybind.cc b/paddle/framework/pybind.cc
index cbb86c4195a..d4ac8fda541 100644
--- a/paddle/framework/pybind.cc
+++ b/paddle/framework/pybind.cc
@@ -200,6 +200,8 @@ All parameter, weight, gradient are variables in Paddle.
     return OpRegistry::CreateOp(desc);
   });
 
+  operator_base.def_static("support_gpu", &OpRegistry::SupportGPU);
+
   operator_base.def("backward",
                     [](const OperatorBase &forwardOp,
                        const std::unordered_set<std::string> &no_grad_vars) {
diff --git a/python/paddle/v2/framework/tests/op_test_util.py b/python/paddle/v2/framework/tests/op_test_util.py
index e6bc7d8a9b5..636828064f7 100644
--- a/python/paddle/v2/framework/tests/op_test_util.py
+++ b/python/paddle/v2/framework/tests/op_test_util.py
@@ -28,7 +28,7 @@ class OpTestMeta(type):
             kwargs = dict()
             places = []
             places.append(core.CPUPlace())
-            if core.is_compile_gpu():
+            if core.is_compile_gpu() and core.Operator.support_gpu(self.type):
                 places.append(core.GPUPlace(0))
 
             for place in places:
@@ -66,7 +66,9 @@ class OpTestMeta(type):
                 for out_name in func.all_output_args:
                     actual = numpy.array(scope.find_var(out_name).get_tensor())
                     expect = self.outputs[out_name]
-                    numpy.isclose(actual, expect)
+                    self.assertTrue(
+                        numpy.allclose(actual, expect),
+                        "output name: " + out_name + "has diff")
 
         obj.test_all = test_all
         return obj
-- 
GitLab