From a37038880e03bc0a985b547bae78f1d0a5d7b048 Mon Sep 17 00:00:00 2001
From: tensor-tang <tangjian03@baidu.com>
Date: Wed, 12 Dec 2018 12:34:08 +0000
Subject: [PATCH] fix unit test with double type

---
 paddle/fluid/operators/jit/helper.h | 44 ++++++++++++++++++-----------
 paddle/fluid/operators/jit/test.cc  | 13 ++++-----
 2 files changed, 33 insertions(+), 24 deletions(-)

diff --git a/paddle/fluid/operators/jit/helper.h b/paddle/fluid/operators/jit/helper.h
index 124587b14..053e5ed07 100644
--- a/paddle/fluid/operators/jit/helper.h
+++ b/paddle/fluid/operators/jit/helper.h
@@ -33,8 +33,11 @@ namespace jit {
 #define EXP_MAX_INPUT 40.0
 
 template <KernelType KT, typename KernelTuples, typename PlaceType>
-inline typename KernelTuples::func_type GetJitCode(
-    typename KernelTuples::attr_type attr) {
+inline typename std::enable_if<
+    std::is_same<typename KernelTuples::data_type, float>::value &&
+        std::is_same<PlaceType, platform::CPUPlace>::value,
+    typename KernelTuples::func_type>::type
+GetJitCode(typename KernelTuples::attr_type attr) {
   using Func = typename KernelTuples::func_type;
   using Attr = typename KernelTuples::attr_type;
   size_t key = JitCodeKey<Attr>(attr);
@@ -45,21 +48,19 @@ inline typename KernelTuples::func_type GetJitCode(
 
   // creator is not related with attr, so can use KernelKey as key
   KernelKey kkey(KT, PlaceType());
-  if (std::is_same<PlaceType, platform::CPUPlace>::value) {
-    // pool: (KernelKey(type, place), vector<GenCreatorPtr>)
-    auto& creator_map = JitCodeCreatorPool().Instance().AllCreators();
-    auto iter = creator_map.find(kkey);
-    if (iter != creator_map.end()) {
-      auto& creators = iter->second;
-      for (auto& cur : creators) {
-        auto i = dynamic_cast<const JitCodeCreator<Attr>*>(cur.get());
-        if (i && i->UseMe(attr)) {
-          auto p = i->CreateJitCode(attr);
-          if (p) {
-            auto f = p->template getCode<Func>();
-            codes.Insert(key, std::move(p));
-            return f;
-          }
+  // pool: (KernelKey(type, place), vector<GenCreatorPtr>)
+  auto& creator_map = JitCodeCreatorPool().Instance().AllCreators();
+  auto iter = creator_map.find(kkey);
+  if (iter != creator_map.end()) {
+    auto& creators = iter->second;
+    for (auto& cur : creators) {
+      auto i = dynamic_cast<const JitCodeCreator<Attr>*>(cur.get());
+      if (i && i->UseMe(attr)) {
+        auto p = i->CreateJitCode(attr);
+        if (p) {
+          auto f = p->template getCode<Func>();
+          codes.Insert(key, std::move(p));
+          return f;
         }
       }
     }
@@ -67,6 +68,15 @@ inline typename KernelTuples::func_type GetJitCode(
   return nullptr;
 }
 
+template <KernelType KT, typename KernelTuples, typename PlaceType>
+inline typename std::enable_if<
+    !std::is_same<typename KernelTuples::data_type, float>::value ||
+        !std::is_same<PlaceType, platform::CPUPlace>::value,
+    typename KernelTuples::func_type>::type
+GetJitCode(typename KernelTuples::attr_type attr) {
+  return nullptr;
+}
+
 // Refer code do not related with attr, which is just for cast
 // Refer is always on CPUPlace
 template <KernelType KT, typename KernelTuples>
diff --git a/paddle/fluid/operators/jit/test.cc b/paddle/fluid/operators/jit/test.cc
index dba7e754e..9ceca2407 100644
--- a/paddle/fluid/operators/jit/test.cc
+++ b/paddle/fluid/operators/jit/test.cc
@@ -48,13 +48,13 @@ void ExpectEQ(const T* target, const T* refer, int n) {
 
 std::vector<int> TestSizes() {
   std::vector<int> s;
-  for (int i = 1; i < 10; ++i) {
+  for (int i = 1; i < 32; ++i) {
     s.push_back(i);
   }
-  // // test some large size
-  // s.push_back(100);
-  // s.push_back(1000);
-  // s.push_back(2000);
+  // test some large size
+  s.push_back(100);
+  s.push_back(1000);
+  s.push_back(2000);
   return s;
 }
 
@@ -148,8 +148,7 @@ void TestXYZNKernel() {
 TEST(JITKernel, vmul) {
   namespace jit = paddle::operators::jit;
   TestXYZNKernel<jit::vmul, float, paddle::platform::CPUPlace>();
-  // TODO(TJ): fix double issue
-  // TestXYZNKernel<jit::vmul, double, paddle::platform::CPUPlace>();
+  TestXYZNKernel<jit::vmul, double, paddle::platform::CPUPlace>();
 }
 
 TEST(JITKernel, vadd) {
-- 
GitLab