Merge pull request #14321 from tensor-tang/fea/jit/vscal

Fea jitcode vscal vaddbias

Author: tensor-tang

paddle/fluid/operators/math/jit_code.cc

@@ -24,21 +24,30 @@ namespace gen {
 
 using namespace platform::jit;  // NOLINT
 
-bool VVVJitCode::init(int d) {
+bool VXXJitCode::init(int d, int scalar_index) {
   // It's not necessary to use avx512 since it would slow down the frequency
   // and this kernel is not compute bound.
-  return MayIUse(avx);
+  return MayIUse(avx) && scalar_index >= 0 && scalar_index <= 2;
 }
 
-void VVVJitCode::generate() {
+void VXXJitCode::generate() {
   // do not need push stack, and do not need save avx512reg if do not use avx512
   int offset = 0;
   if (with_relu_) {
     vxorps(ymm_zero, ymm_zero, ymm_zero);
   }
+  if (scalar_index_ == 1) {
+    vbroadcastss(ymm_src1, ptr[param1]);
+  } else if (scalar_index_ == 2) {
+    vbroadcastss(ymm_src2, ptr[param2]);
+  }
   for (int i = 0; i < num_ / AVX_FLOAT_BLOCK; ++i) {
-    vmovups(ymm_src1, ptr[param1 + offset]);
-    vmovups(ymm_src2, ptr[param2 + offset]);
+    if (scalar_index_ != 1) {
+      vmovups(ymm_src1, ptr[param1 + offset]);
+    }
+    if (scalar_index_ != 2) {
+      vmovups(ymm_src2, ptr[param2 + offset]);
+    }
     if (type_ == operand_type::mul) {
       vmulps(ymm_dst, ymm_src1, ymm_src2);
     } else if (type_ == operand_type::add) {
@@ -52,8 +61,12 @@ void VVVJitCode::generate() {
   }
   int rest = num_ % AVX_FLOAT_BLOCK;
   if (rest >= 4) {
-    vmovups(xmm_src1, ptr[param1 + offset]);
-    vmovups(xmm_src2, ptr[param2 + offset]);
+    if (scalar_index_ != 1) {
+      vmovups(xmm_src1, ptr[param1 + offset]);
+    }
+    if (scalar_index_ != 2) {
+      vmovups(xmm_src2, ptr[param2 + offset]);
+    }
     if (type_ == operand_type::mul) {
       vmulps(xmm_dst, xmm_src1, xmm_src2);
     } else if (type_ == operand_type::add) {
@@ -67,8 +80,12 @@ void VVVJitCode::generate() {
     rest -= 4;
   }
   if (rest >= 2) {
-    vmovq(xmm_src1, ptr[param1 + offset]);
-    vmovq(xmm_src2, ptr[param2 + offset]);
+    if (scalar_index_ != 1) {
+      vmovups(xmm_src1, ptr[param1 + offset]);
+    }
+    if (scalar_index_ != 2) {
+      vmovups(xmm_src2, ptr[param2 + offset]);
+    }
     if (type_ == operand_type::mul) {
       vmulps(xmm_dst, xmm_src1, xmm_src2);
     } else if (type_ == operand_type::add) {
@@ -82,8 +99,12 @@ void VVVJitCode::generate() {
     rest -= 2;
   }
   if (rest > 0) {
-    vmovss(xmm_src1, ptr[param1 + offset]);
-    vmovss(xmm_src2, ptr[param2 + offset]);
+    if (scalar_index_ != 1) {
+      vmovups(xmm_src1, ptr[param1 + offset]);
+    }
+    if (scalar_index_ != 2) {
+      vmovups(xmm_src2, ptr[param2 + offset]);
+    }
     if (type_ == operand_type::mul) {
       vmulss(xmm_dst, xmm_src1, xmm_src2);
     } else if (type_ == operand_type::add) {
@@ -96,6 +117,7 @@ void VVVJitCode::generate() {
   }
   ret();
 }
+
 }  // namespace gen
 }  // namespace jitkernel
 }  // namespace math

paddle/fluid/operators/math/jit_code.h

@@ -29,47 +29,60 @@ using ymm_t = const Xbyak::Ymm;
 using zmm_t = const Xbyak::Zmm;
 using Label = Xbyak::Label;
 
-// function: vec = Operand(vec, vec) (maybe with relu)
 typedef enum { mul = 0, add } operand_type;
 
-class VVVJitCode : public JitCode {
+// function: vec = Operand(vec(or scalar), vec(or scalar)) (maybe with relu)
+class VXXJitCode : public JitCode {
  public:
   const char* name() const override {
-    std::string base = "VVVJitCode";
+    std::string base = "VXXJitCode";
+    if (scalar_index_ == 1) {
+      base += "_Scalar";
+    } else {
+      base += "_Vec";
+    }
     if (type_ == operand_type::mul) {
       base += "_Mul";
     } else if (type_ == operand_type::add) {
       base += "_Add";
     }
-    base += (with_relu_ ? "_relu" : "");
+    if (scalar_index_ == 2) {
+      base += "_Scalar";
+    } else {
+      base += "_Vec";
+    }
+    base += (with_relu_ ? "_Relu" : "");
     return base.c_str();
   }
-  explicit VVVJitCode(int d, operand_type type, bool with_relu,
-                      size_t code_size = 256 * 1024, void* code_ptr = nullptr)
+  explicit VXXJitCode(int d, operand_type type, int scalar_index,
+                      bool with_relu, size_t code_size = 256 * 1024,
+                      void* code_ptr = nullptr)
       : JitCode(code_size, code_ptr),
         num_(d),
         type_(type),
+        scalar_index_(scalar_index),
         with_relu_(with_relu) {}
-  static bool init(int d);
+  static bool init(int d, int scalar_index = 0);
   void generate() override;
 
  private:
   int num_;
   operand_type type_;
+  int scalar_index_;
   bool with_relu_;
   reg64_t param1{abi_param1};
   reg64_t param2{abi_param2};
   reg64_t param3{abi_param3};
 
   xmm_t xmm_src1 = xmm_t(0);
   xmm_t xmm_src2 = xmm_t(1);
-  xmm_t xmm_dst = xmm_t(1);
-  xmm_t xmm_zero = xmm_t(2);
+  xmm_t xmm_dst = xmm_t(2);
+  xmm_t xmm_zero = xmm_t(3);
 
   ymm_t ymm_src1 = ymm_t(0);
   ymm_t ymm_src2 = ymm_t(1);
-  ymm_t ymm_dst = ymm_t(1);
-  ymm_t ymm_zero = ymm_t(2);
+  ymm_t ymm_dst = ymm_t(2);
+  ymm_t ymm_zero = ymm_t(3);
 };
 
 }  // namespace gen

paddle/fluid/operators/math/jit_kernel.h

@@ -83,14 +83,15 @@ class VAddReluKernel : public Kernel {
 template <typename T>
 class VScalKernel : public Kernel {
  public:
-  virtual void Compute(const T a, const T *x, T *y) const = 0;
-  virtual void Compute(const T a, T *x) const = 0;
+  // y = a.*x
+  void (*Compute)(const T *, const T *, T *, int);
 };
 
 template <typename T>
 class VAddBiasKernel : public Kernel {
  public:
-  virtual void Compute(const T a, const T *x, T *y) const = 0;
+  // y = a.+x
+  void (*Compute)(const T *, const T *, T *, int);
 };
 
 template <typename T>