From 9255119fd915e1ec58ae60d18f3012305383d8f9 Mon Sep 17 00:00:00 2001
From: tensor-tang <tangjian03@baidu.com>
Date: Mon, 5 Nov 2018 06:09:09 +0000
Subject: [PATCH] refine jit vmul with all size

---
 paddle/fluid/operators/math/jit_code.cc | 21 ++++++++++-----------
 paddle/fluid/operators/math/jit_code.h  |  1 -
 2 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/paddle/fluid/operators/math/jit_code.cc b/paddle/fluid/operators/math/jit_code.cc
index c3bb60f2a8..9e2cc18c7a 100644
--- a/paddle/fluid/operators/math/jit_code.cc
+++ b/paddle/fluid/operators/math/jit_code.cc
@@ -27,11 +27,7 @@ using namespace platform::jit;  // NOLINT
 bool VMulJitCode::init(int d) {
   // It's not necessary to use avx512 since it would slow down the frequency
   // and this kernel is not compute bound.
-  if (MayIUse(avx)) {
-    return d % 2 == 0;
-  } else {
-    return false;
-  }
+  return MayIUse(avx);
 }
 
 void VMulJitCode::generate() {
@@ -54,16 +50,19 @@ void VMulJitCode::generate() {
     rest -= 4;
   }
   if (rest >= 2) {
-    mov(tmp, qword[param1 + offset]);
-    vmovq(xmm_src1, tmp);
-    mov(tmp, qword[param2 + offset]);
-    vmovq(xmm_src2, tmp);
+    vmovq(xmm_src1, ptr[param1 + offset]);
+    vmovq(xmm_src2, ptr[param2 + offset]);
     vmulps(xmm_dst, xmm_src1, xmm_src2);
-    vmovq(tmp, xmm_dst);
-    mov(ptr[param3 + offset], tmp);
+    vmovq(ptr[param3 + offset], xmm_dst);
     offset += sizeof(float) * 2;
     rest -= 2;
   }
+  if (rest > 0) {
+    vmovss(xmm_src1, ptr[param1 + offset]);
+    vmovss(xmm_src2, ptr[param2 + offset]);
+    vmulss(xmm_dst, xmm_src1, xmm_src2);
+    vmovss(ptr[param3 + offset], xmm_dst);
+  }
   ret();
 }
 
diff --git a/paddle/fluid/operators/math/jit_code.h b/paddle/fluid/operators/math/jit_code.h
index c77252a326..6007b29081 100644
--- a/paddle/fluid/operators/math/jit_code.h
+++ b/paddle/fluid/operators/math/jit_code.h
@@ -43,7 +43,6 @@ class VMulJitCode : public JitCode {
   reg64_t param1{abi_param1};
   reg64_t param2{abi_param2};
   reg64_t param3{abi_param3};
-  reg64_t tmp = rax;
 
   xmm_t xmm_src1 = xmm_t(0);
   xmm_t xmm_src2 = xmm_t(1);
-- 
GitLab