From ce674b685f89a66ef0b9163b76b456a9580e5f41 Mon Sep 17 00:00:00 2001
From: tensor-tang <tangjian03@baidu.com>
Date: Thu, 6 Dec 2018 05:24:50 +0000
Subject: [PATCH] add readme doc and complete TODOs

---
 paddle/fluid/operators/jitkernels/README.md   | 46 ++++++++++++++++++-
 .../fluid/operators/jitkernels/jitcode_base.h |  6 +--
 .../fluid/operators/jitkernels/kernel_base.h  | 13 ++++--
 .../fluid/operators/jitkernels/kernel_pool.h  | 14 +-----
 .../fluid/operators/jitkernels/more/mkl/mkl.h | 11 +----
 .../fluid/operators/jitkernels/refer/refer.h  |  4 +-
 paddle/fluid/operators/jitkernels/registry.h  |  2 +-
 paddle/fluid/operators/jitkernels/test.cc     |  8 ++--
 8 files changed, 65 insertions(+), 39 deletions(-)

diff --git a/paddle/fluid/operators/jitkernels/README.md b/paddle/fluid/operators/jitkernels/README.md
index 3401e9be531..fd6428b43ec 100644
--- a/paddle/fluid/operators/jitkernels/README.md
+++ b/paddle/fluid/operators/jitkernels/README.md
@@ -1,4 +1,46 @@
-TBD
+# JIT Kernel
+
+结合函数模板和JIT生成需要的kernel函数。
+这里的kernel是比Operator中kernel更小级别的算子单元，更侧重的是在不同硬件上的性能。
+目前仅支持CPU上的高性能计算。
+
+## 目录结构
+
+```txt
+PaddlePaddle/Paddle/paddle/fluid/
+├── ...
+├── operator/
+│   ├── .../
+└── jit/
+    ├── ...
+    ├── jitcode/
+    │   └── ...
+    |── more/
+    │   ├── ...
+    │   ├── mkl/
+    │   │   └── ...
+    │   └── openblas/
+    │       └── ...
+    └── refer/
+        └── ...
+```
+
+基础class都的根目录下，根目录下包括jitcode,more和refer。每个目录下都是一种实现，每种kernel算子都需要有reference的实现，其他的都是可选的。
+- jitcode： 代表使用jit生成的code，需要依赖xbyak。他关心的是性能。
+- refer：代表reference的实现，每种kernel算子都需要有在CPU上的reference的实现，他主要关心的算法逻辑。
+- more： 下面可以放入跟多实现，包括mkl，mkldnn，openblas等，也可以是自身已有的kernel组合。
 
-# Use me
+## 动态获取
+
+提供一个get方法，根据kernel类别获取，每种实现都有自己的使用范围，根据范围动态和当前条件选择需要的kernel函数。
+
+## 测试
+
+- 逻辑测试
+    所有实现都要与refer的code对比，需要满足精度要求
+- 性能测试
+
+# 如何添加新的算子
+TBD
+## Use me
 Add USE_JIT_KERNEL(yourname) to CMakefile.
diff --git a/paddle/fluid/operators/jitkernels/jitcode_base.h b/paddle/fluid/operators/jitkernels/jitcode_base.h
index ffec62163a7..de8aaf229fe 100644
--- a/paddle/fluid/operators/jitkernels/jitcode_base.h
+++ b/paddle/fluid/operators/jitkernels/jitcode_base.h
@@ -62,11 +62,7 @@ class JitBase : public Kernel {
 };
 
 template <KernelType KT, typename T, typename Attr>
-std::unique_ptr<JitBase> CreateJitCode(Attr attr);  //{
-//   if (UseJitCode<KT,T,Attr>) {
-//     return make_unique<xxxxclass>(attr, CodeSize<KT,T,Attr>());
-//   }
-// }
+std::unique_ptr<JitBase> CreateJitCode(Attr attr);
 
 }  // namespace jitkernels
 }  // namespace operators
diff --git a/paddle/fluid/operators/jitkernels/kernel_base.h b/paddle/fluid/operators/jitkernels/kernel_base.h
index eeaa0617cb8..6fbb0f9f7ea 100644
--- a/paddle/fluid/operators/jitkernels/kernel_base.h
+++ b/paddle/fluid/operators/jitkernels/kernel_base.h
@@ -21,6 +21,13 @@ namespace jitkernels {
 
 typedef enum { vmul = 0, vadd = 1, vsub, vexp } KernelType;
 
+template <typename T>
+struct VMulTypes {
+  typedef T data_type;
+  typedef int attr_type;
+  typedef void (*func_type)(const T*, const T*, T*, int);
+};
+
 // Just for adding to kernel pool without template
 class Kernel {
  public:
@@ -29,10 +36,10 @@ class Kernel {
   DISABLE_COPY_AND_ASSIGN(Kernel);
 };
 
-template <typename T, typename Func, typename Attr>  // TODO(TJ): use tuple
+template <typename T, typename Func, typename Attr>
 class KernelImpl : public Kernel {
  public:
-  using ELEMENT_TYPE = T;  // TODO(TJ): remove me?
+  using ELEMENT_TYPE = T;
   virtual Func GetFunc() const { return func; }
   virtual bool UseMe(Attr attr) const = 0;
 
@@ -40,7 +47,7 @@ class KernelImpl : public Kernel {
   Func func{nullptr};
 };
 
-template <typename T, typename Func, typename Attr>  // TODO(TJ): use tuple
+template <typename T, typename Func, typename Attr>
 class ReferKernel : public KernelImpl<T, Func, Attr> {
  public:
   // Refer code can always be used
diff --git a/paddle/fluid/operators/jitkernels/kernel_pool.h b/paddle/fluid/operators/jitkernels/kernel_pool.h
index f398093dfe2..901a891cb38 100644
--- a/paddle/fluid/operators/jitkernels/kernel_pool.h
+++ b/paddle/fluid/operators/jitkernels/kernel_pool.h
@@ -27,8 +27,6 @@ namespace paddle {
 namespace operators {
 namespace jitkernels {
 
-// TODO(TJ): rename file to kernel_pool
-
 template <KernelType KT>
 class JitCodePool {
   typedef std::unique_ptr<JitBase> JitBasePtr;
@@ -54,14 +52,6 @@ class JitCodePool {
   DISABLE_COPY_AND_ASSIGN(JitCodePool);
 };
 
-// TODO(TJ): std::tuple<T, Func, Attr>
-// template <typename T, typename Func, typename Attr>
-// struct KernelAttr {
-//   typedef T data_type;
-//   typedef Func return_type;
-//   typedef Attr attr_type;
-// };
-
 typedef std::unique_ptr<const Kernel> KernelPtr;
 typedef std::unordered_map<KernelKey, std::vector<KernelPtr>, KernelKey::Hash>
     KernelMap;
@@ -120,7 +110,6 @@ inline Func GetRefer() {
   return nullptr;
 }
 
-// TODO(TJ): make tuple? named KernelAttr
 template <KernelType KT, typename T, typename Func, typename Attr,
           typename PlaceType = platform::CPUPlace>
 const Func Get(Attr attr) {
@@ -130,8 +119,7 @@ const Func Get(Attr attr) {
     return codes.AllKernels().at(key)->template getCode<Func>();
   }
 
-  if (std::is_same<PlaceType, platform::CPUPlace>::value) {  // TODO(TJ): float
-                                                             // move to create
+  if (std::is_same<PlaceType, platform::CPUPlace>::value) {
     auto p = CreateJitCode<KT, T, Attr>(attr);
     if (p) {
       auto f = p->template getCode<Func>();
diff --git a/paddle/fluid/operators/jitkernels/more/mkl/mkl.h b/paddle/fluid/operators/jitkernels/more/mkl/mkl.h
index 75ed34ef48e..9cf032db43f 100644
--- a/paddle/fluid/operators/jitkernels/more/mkl/mkl.h
+++ b/paddle/fluid/operators/jitkernels/more/mkl/mkl.h
@@ -27,16 +27,9 @@ namespace mkl {
 template <typename T>
 void VMul(const T* x, const T* y, T* z, int n);
 
-// template <typename T>
-// struct VMulTypes{
-// typedef T date_type;
-// typedef void (*func)(const T*, const T*, T*, int) func_type;
-//   typedef int attr_type;
-// };
-
 template <typename T>
-class VMulKernel
-    : public KernelImpl<T, void (*)(const T*, const T*, T*, int), int> {
+class VMulKernel : public KernelImpl<T, typename VMulTypes<T>::func_type,
+                                     typename VMulTypes<T>::attr_type> {
  public:
   VMulKernel() { this->func = VMul<T>; }
   bool UseMe(int d) const override {
diff --git a/paddle/fluid/operators/jitkernels/refer/refer.h b/paddle/fluid/operators/jitkernels/refer/refer.h
index 163c6d73dce..796f58d4017 100644
--- a/paddle/fluid/operators/jitkernels/refer/refer.h
+++ b/paddle/fluid/operators/jitkernels/refer/refer.h
@@ -29,8 +29,8 @@ void VMul(const T* x, const T* y, T* z, int n) {
 }
 
 template <typename T>
-class VMulKernel
-    : public ReferKernel<T, void (*)(const T*, const T*, T*, int), int> {
+class VMulKernel : public ReferKernel<T, typename VMulTypes<T>::func_type,
+                                      typename VMulTypes<T>::attr_type> {
  public:
   VMulKernel() { this->func = VMul<T>; }
 };
diff --git a/paddle/fluid/operators/jitkernels/registry.h b/paddle/fluid/operators/jitkernels/registry.h
index cd414bb096c..6d817461bec 100644
--- a/paddle/fluid/operators/jitkernels/registry.h
+++ b/paddle/fluid/operators/jitkernels/registry.h
@@ -26,7 +26,7 @@ namespace paddle {
 namespace operators {
 namespace jitkernels {
 
-// make_unique is supported from c++14
+// make_unique is supported since c++14
 template <typename T, typename... Args>
 inline std::unique_ptr<T> make_unique(Args&&... args) {
   static_assert(!std::is_array<T>::value, "T must not be array");
diff --git a/paddle/fluid/operators/jitkernels/test.cc b/paddle/fluid/operators/jitkernels/test.cc
index eb0d30eecdb..d27b5d1cbae 100644
--- a/paddle/fluid/operators/jitkernels/test.cc
+++ b/paddle/fluid/operators/jitkernels/test.cc
@@ -69,10 +69,10 @@ TEST(JitKernel, vmul) {
   namespace jit = paddle::operators::jitkernels;
   // TODO(TJ): test more vector size
   for (int d = 1; d < 30; ++d) {
-    auto ref = jit::GetRefer<jit::vmul, T,
-                             void (*)(const T*, const T*, T*, int), int>();
-    auto tgt = jit::Get<jit::vmul, T, void (*)(const T*, const T*, T*, int),
-                        int, PlaceType>(d);
+    auto ref = jit::GetRefer<jit::vmul, T, jit::VMulTypes<T>::func_type,
+                             jit::VMulTypes<T>::attr_type>();
+    auto tgt = jit::Get<jit::vmul, T, jit::VMulTypes<T>::func_type,
+                        jit::VMulTypes<T>::attr_type, PlaceType>(d);
     EXPECT_TRUE(ref != nullptr);
     EXPECT_TRUE(tgt != nullptr);
 
-- 
GitLab