add refer vrelu, videntity, vexp, vsigmoid, vtanh and test and benchmark

paddle/fluid/operators/jit/benchmark.cc

@@ -201,6 +201,77 @@ void BenchAXYNKernel() {
   }
 }
 
+// return this function avg time
+template <typename T, typename KernelTuples>
+double BenchXYNFunc(const typename KernelTuples::func_type tgt,
+                    const std::vector<T>& x,
+                    std::vector<T>& y) {  // NOLINT
+  const T* x_data = x.data();
+  T* y_data = y.data();
+  const int d = y.size();
+  for (int i = 0; i < FLAGS_burning; ++i) {
+    tgt(x_data, y_data, d);
+  }
+  auto start = GetCurrentUS();
+  for (int i = 0; i < FLAGS_repeat; ++i) {
+    tgt(x_data, y_data, d);
+  }
+  auto end = GetCurrentUS();
+  return (end - start) / FLAGS_repeat;
+}
+
+template <paddle::operators::jit::KernelType KT, typename T, typename PlaceType>
+void BenchXYNKernel() {
+  namespace jit = paddle::operators::jit;
+  for (int d : TestSizes()) {
+    std::vector<std::pair<std::string, double>> infos;
+    std::vector<T> x(d), y(d);
+    RandomVec<T>(d, x.data());
+    // test refer
+    auto refer = jit::GetRefer<KT, jit::XYNTuples<T>>();
+    if (refer) {
+      auto res = BenchXYNFunc<T, jit::XYNTuples<T>>(refer, x, y);
+      infos.push_back(std::make_pair("Refer", res));
+    }
+    // test jitcode
+    auto jitcode = jit::GetJitCode<KT, jit::XYNTuples<T>, PlaceType>(d);
+    if (jitcode) {
+      auto res = BenchXYNFunc<T, jit::XYNTuples<T>>(jitcode, x, y);
+      infos.push_back(std::make_pair("JitCode", res));
+    }
+    // test all impls in more
+    jit::KernelKey kkey(KT, PlaceType());
+    auto& pool = jit::KernelPool().Instance().AllKernels();
+    auto iter = pool.find(kkey);
+    if (iter != pool.end()) {
+      auto& impls = iter->second;
+      for (auto& impl : impls) {
+        auto i =
+            dynamic_cast<const jit::KernelImpl<jit::XYNTuples<T>>*>(impl.get());
+        if (i && i->UseMe(d)) {
+          auto more = i->GetFunc();
+          auto res = BenchXYNFunc<T, jit::XYNTuples<T>>(more, x, y);
+          infos.push_back(std::make_pair("More", res));
+        }
+      }
+    }
+    // Test result from Get function
+    auto tgt = jit::Get<KT, jit::XYNTuples<T>, PlaceType>(d);
+    if (!tgt) {
+      LOG(ERROR) << "Target can not be empty!";
+    }
+    auto res = BenchXYNFunc<T, jit::XYNTuples<T>>(tgt, x, y);
+    infos.push_back(std::make_pair("Target", res));
+    // print
+    std::ostringstream loginfos;
+    loginfos << "Kernel Type: " << jit::to_string(KT) << ", size " << d << ": ";
+    for (auto pair : infos) {
+      loginfos << pair.first << " takes " << pair.second << " us; ";
+    }
+    LOG(INFO) << loginfos.str();
+  }
+}
+
 // Benchmark all jit kernels including jitcode, mkl and refer.
 // To use this tool, run command: ./benchmark [options...]
 // Options:
@@ -222,4 +293,10 @@ int main(int argc, char* argv[]) {
 
   BenchAXYNKernel<jit::vscal, T, PlaceType>();
   BenchAXYNKernel<jit::vaddbias, T, PlaceType>();
+
+  BenchXYNKernel<jit::vrelu, T, PlaceType>();
+  BenchXYNKernel<jit::videntity, T, PlaceType>();
+  BenchXYNKernel<jit::vexp, T, PlaceType>();
+  BenchXYNKernel<jit::vsigmoid, T, PlaceType>();
+  BenchXYNKernel<jit::vtanh, T, PlaceType>();
 }

paddle/fluid/operators/jit/helper.cc

@@ -19,28 +19,30 @@ namespace paddle {
 namespace operators {
 namespace jit {
 
+#define ONE_CASE(key) \
+  case key:           \
+    return #key
+
 const char* to_string(KernelType kt) {
   switch (kt) {
-    case vmul:
-      return "vmul";
-    case vadd:
-      return "vadd";
-    case vaddrelu:
-      return "vaddrelu";
-    case vsub:
-      return "vsub";
-    case vscal:
-      return "vscal";
-    case vexp:
-      return "vexp";
-    case vaddbias:
-      return "vaddbias";
+    ONE_CASE(vmul);
+    ONE_CASE(vadd);
+    ONE_CASE(vaddrelu);
+    ONE_CASE(vsub);
+    ONE_CASE(vscal);
+    ONE_CASE(vaddbias);
+    ONE_CASE(vrelu);
+    ONE_CASE(videntity);
+    ONE_CASE(vexp);
+    ONE_CASE(vsigmoid);
+    ONE_CASE(vtanh);
     default:
       PADDLE_THROW("Not support type: %d", kt);
       return "NOT JITKernel";
   }
   return nullptr;
 }
+#undef ONE_CASE
 
 }  // namespace jit
 }  // namespace operators

paddle/fluid/operators/jit/kernel_base.h

@@ -26,7 +26,11 @@ typedef enum {
   vsub,
   vscal,
   vaddbias,
-  vexp
+  vrelu,
+  videntity,
+  vexp,
+  vsigmoid,
+  vtanh
 } KernelType;
 
 template <typename T>
@@ -39,6 +43,13 @@ struct XYZNTuples {
 template <typename T>
 struct AXYNTuples : public XYZNTuples<T> {};
 
+template <typename T>
+struct XYNTuples {
+  typedef T data_type;
+  typedef int attr_type;
+  typedef void (*func_type)(const T*, T*, int);
+};
+
 // Just for adding to kernel pool without template
 class Kernel {
  public:

paddle/fluid/operators/jit/refer/CMakeLists.txt

@@ -13,3 +13,8 @@ USE_JITKERNEL_REFER(vaddrelu)
 USE_JITKERNEL_REFER(vsub)
 USE_JITKERNEL_REFER(vscal)
 USE_JITKERNEL_REFER(vaddbias)
+USE_JITKERNEL_REFER(vrelu)
+USE_JITKERNEL_REFER(videntity)
+USE_JITKERNEL_REFER(vexp)
+USE_JITKERNEL_REFER(vsigmoid)
+USE_JITKERNEL_REFER(vtanh)

paddle/fluid/operators/jit/refer/refer.cc

@@ -29,4 +29,10 @@ REGISTER_REFER_KERNEL(vsub, VSub);
 REGISTER_REFER_KERNEL(vscal, VScal);
 REGISTER_REFER_KERNEL(vaddbias, VAddBias);
 
+REGISTER_REFER_KERNEL(vrelu, VRelu);
+REGISTER_REFER_KERNEL(videntity, VIdentity);
+REGISTER_REFER_KERNEL(vexp, VExp);
+REGISTER_REFER_KERNEL(vsigmoid, VSigmoid);
+REGISTER_REFER_KERNEL(vtanh, VTanh);
+
 #undef REGISTER_REFER_KERNEL

paddle/fluid/operators/jit/refer/refer.h

@@ -66,6 +66,50 @@ void VAddBias(const T* a, const T* x, T* y, int n) {
   }
 }
 
+template <typename T>
+void VRelu(const T* x, T* y, int n) {
+  for (int i = 0; i < n; ++i) {
+    y[i] = x[i] > 0 ? x[i] : 0;
+  }
+}
+
+template <typename T>
+inline void VIdentity(const T* x, T* y, int n) {
+  for (int i = 0; i < n; ++i) {
+    y[i] = x[i];
+  }
+}
+
+template <typename T>
+void VExp(const T* x, T* y, int n) {
+  for (int i = 0; i < n; ++i) {
+    y[i] = std::exp(x[i]);
+  }
+}
+
+template <typename T>
+void VSigmoid(const T* x, T* y, int n) {
+  // y = 1 / (1 + e^-x)
+  const T min = SIGMOID_THRESHOLD_MIN;
+  const T max = SIGMOID_THRESHOLD_MAX;
+  for (int i = 0; i < n; ++i) {
+    T tmp = (x[i] < min) ? min : ((x[i] > max) ? max : x[i]);
+    y[i] = static_cast<T>(1) / (static_cast<T>(1) + std::exp(-tmp));
+  }
+}
+
+template <typename T>
+void VTanh(const T* x, T* y, int n) {
+  // y = 2 * sigmoid(2x) - 1
+  for (int i = 0; i < n; ++i) {
+    y[i] = static_cast<T>(2) * x[i];
+  }
+  VSigmoid(y, y, n);
+  for (int i = 0; i < n; ++i) {
+    y[i] = static_cast<T>(2) * y[i] - static_cast<T>(1);
+  }
+}
+
 #define DECLARE_REFER_KERNEL(name, tuples)             \
   template <typename T>                                \
   class name##Kernel : public ReferKernel<tuples<T>> { \
@@ -83,6 +127,13 @@ DECLARE_REFER_KERNEL(VSub, XYZNTuples);
 DECLARE_REFER_KERNEL(VScal, AXYNTuples);
 DECLARE_REFER_KERNEL(VAddBias, AXYNTuples);
 
+// const T* x, T* y, int n
+DECLARE_REFER_KERNEL(VRelu, XYNTuples);
+DECLARE_REFER_KERNEL(VIdentity, XYNTuples);
+DECLARE_REFER_KERNEL(VExp, XYNTuples);
+DECLARE_REFER_KERNEL(VSigmoid, XYNTuples);
+DECLARE_REFER_KERNEL(VTanh, XYNTuples);
+
 #undef DECLARE_REFER_KERNEL
 
 }  // namespace refer

paddle/fluid/operators/jit/test.cc

@@ -250,6 +250,106 @@ TEST(JITKernel, vaddbias) {
   TestAXYNKernel<jit::vaddbias, double, paddle::platform::CPUPlace>();
 }
 
+template <typename T, typename KernelTuples>
+void TestXYNFunc(const typename KernelTuples::func_type tgt,
+                 const std::vector<T>& x, const std::vector<T>& yref) {
+  EXPECT_TRUE(tgt != nullptr);
+  EXPECT_EQ(yref.size(), x.size());
+  const T* x_data = x.data();
+  const T* yref_data = yref.data();
+  const int d = yref.size();
+  std::vector<T> ytgt(d);
+  T* ytgt_data = ytgt.data();
+  // test normal
+  tgt(x_data, ytgt_data, d);
+  ExpectEQ<T>(ytgt_data, yref_data, d);
+  // test inplace x
+  std::copy(x.begin(), x.end(), ytgt.begin());
+  tgt(ytgt_data, ytgt_data, d);
+  ExpectEQ<T>(ytgt_data, yref_data, d);
+}
+
+template <paddle::operators::jit::KernelType KT, typename T, typename PlaceType>
+void TestXYNKernel() {
+  namespace jit = paddle::operators::jit;
+  VLOG(10) << "===== Test JITKernel " << jit::to_string(KT);
+  for (int d : TestSizes()) {
+    auto ref = jit::GetRefer<KT, jit::XYNTuples<T>>();
+    EXPECT_TRUE(ref != nullptr);
+
+    std::vector<T> x(d), yref(d);
+    std::vector<T> xinp(d);  // inplace test
+    RandomVec<T>(d, x.data());
+    std::copy(x.begin(), x.end(), xinp.begin());
+
+    const T* x_data = x.data();
+    T* yref_data = yref.data();
+    T* xinp_data = xinp.data();
+    // test refer code inplace
+    ref(x_data, yref_data, d);
+    ref(xinp_data, xinp_data, d);
+    ExpectEQ<T>(xinp_data, yref_data, d);
+
+    // test jitcode
+    auto jitcode = jit::GetJitCode<KT, jit::XYNTuples<T>, PlaceType>(d);
+    if (jitcode) {
+      VLOG(10) << "Test Jitcode Kernel, size: " << d;
+      TestXYNFunc<T, jit::XYNTuples<T>>(jitcode, x, yref);
+    }
+
+    // test all impls in more
+    jit::KernelKey kkey(KT, PlaceType());
+    auto& pool = jit::KernelPool().Instance().AllKernels();
+    auto iter = pool.find(kkey);
+    if (iter != pool.end()) {
+      auto& impls = iter->second;
+      for (auto& impl : impls) {
+        auto i =
+            dynamic_cast<const jit::KernelImpl<jit::XYNTuples<T>>*>(impl.get());
+        if (i && i->UseMe(d)) {
+          auto more = i->GetFunc();
+          VLOG(10) << "Test More Kernel, size: " << d;
+          TestXYNFunc<T, jit::XYNTuples<T>>(more, x, yref);
+        }
+      }
+    }
+    // Test result from Get function
+    VLOG(10) << "Test Get function, size: " << d;
+    auto tgt = jit::Get<KT, jit::XYNTuples<T>, PlaceType>(d);
+    TestXYNFunc<T, jit::XYNTuples<T>>(tgt, x, yref);
+  }
+}
+
+TEST(JITKernel, vrelu) {
+  namespace jit = paddle::operators::jit;
+  TestXYNKernel<jit::vrelu, float, paddle::platform::CPUPlace>();
+  TestXYNKernel<jit::vrelu, double, paddle::platform::CPUPlace>();
+}
+
+TEST(JITKernel, videntity) {
+  namespace jit = paddle::operators::jit;
+  TestXYNKernel<jit::videntity, float, paddle::platform::CPUPlace>();
+  TestXYNKernel<jit::videntity, double, paddle::platform::CPUPlace>();
+}
+
+TEST(JITKernel, vexp) {
+  namespace jit = paddle::operators::jit;
+  TestXYNKernel<jit::vexp, float, paddle::platform::CPUPlace>();
+  TestXYNKernel<jit::vexp, double, paddle::platform::CPUPlace>();
+}
+
+TEST(JITKernel, vsigmoid) {
+  namespace jit = paddle::operators::jit;
+  TestXYNKernel<jit::vsigmoid, float, paddle::platform::CPUPlace>();
+  TestXYNKernel<jit::vsigmoid, double, paddle::platform::CPUPlace>();
+}
+
+TEST(JITKernel, vtanh) {
+  namespace jit = paddle::operators::jit;
+  TestXYNKernel<jit::vtanh, float, paddle::platform::CPUPlace>();
+  TestXYNKernel<jit::vtanh, double, paddle::platform::CPUPlace>();
+}
+
 TEST(JITKernel, pool) {
   // TODO(TJ): add some test
 }

paddle/fluid/operators/math/jit_kernel_refer.h

@@ -24,46 +24,6 @@ namespace math {
 namespace jitkernel {
 namespace refer {
 
-template <typename T>
-void VRelu(const T* x, T* y, int n) {
-  for (int i = 0; i < n; ++i) {
-    y[i] = x[i] > 0 ? x[i] : 0;
-  }
-}
-
-template <typename T>
-inline void VIdentity(const T* x, T* y, int n) {}
-
-template <typename T>
-void VExp(const T* x, T* y, int n) {
-  for (int i = 0; i < n; ++i) {
-    y[i] = std::exp(x[i]);
-  }
-}
-
-template <typename T>
-void VSigmoid(const T* x, T* y, int n) {
-  // y = 1 / (1 + e^-x)
-  const T min = SIGMOID_THRESHOLD_MIN;
-  const T max = SIGMOID_THRESHOLD_MAX;
-  for (int i = 0; i < n; ++i) {
-    T tmp = (x[i] < min) ? min : ((x[i] > max) ? max : x[i]);
-    y[i] = static_cast<T>(1) / (static_cast<T>(1) + std::exp(-tmp));
-  }
-}
-
-template <typename T>
-void VTanh(const T* x, T* y, int n) {
-  // y = 2 * sigmoid(2x) - 1
-  for (int i = 0; i < n; ++i) {
-    y[i] = static_cast<T>(2) * x[i];
-  }
-  VSigmoid(y, y, n);
-  for (int i = 0; i < n; ++i) {
-    y[i] = static_cast<T>(2) * y[i] - static_cast<T>(1);
-  }
-}
-
 template <typename T>
 void (*getActFunc(const std::string& type))(const T*, T*, int) {  // NOLINT
   if (type == "sigmoid") {