jit_kernel_test.cc 26.7 KB
Newer Older
T
tensor-tang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#include "paddle/fluid/operators/math/jit_kernel.h"
T
tensor-tang 已提交
16
#include <sys/time.h>
T
tensor-tang 已提交
17
#include <cmath>    // for exp
T
tensor-tang 已提交
18
#include <cstring>  // for memcpy
T
tensor-tang 已提交
19
#include <random>
T
tensor-tang 已提交
20 21 22 23 24 25
#include <string>
#include <vector>
#include "gflags/gflags.h"
#include "glog/logging.h"
#include "gtest/gtest.h"

T
tensor-tang 已提交
26 27 28 29 30 31 32 33
#ifdef PADDLE_WITH_MKLML
#include "paddle/fluid/platform/dynload/mklml.h"
#endif

#ifdef __AVX__
#include <immintrin.h>
#endif

T
tensor-tang 已提交
34 35
constexpr int repeat = 20000;

T
tensor-tang 已提交
36 37 38
// TODO(TJ): benchmark and test should be seperated,
// benchmark should verify more sizes

T
tensor-tang 已提交
39 40 41 42 43 44 45
inline double GetCurrentUS() {
  struct timeval time;
  gettimeofday(&time, NULL);
  return 1e+6 * time.tv_sec + time.tv_usec;
}

template <typename T>
T
tensor-tang 已提交
46 47
void RandomVec(const int n, T* a, const T lower = static_cast<T>(-20.f),
               const T upper = static_cast<T>(20.f)) {
T
tensor-tang 已提交
48 49 50 51 52 53 54 55
  static unsigned int seed = 100;
  std::mt19937 rng(seed++);
  std::uniform_real_distribution<double> uniform_dist(0, 1);
  for (int i = 0; i < n; ++i) {
    a[i] = static_cast<T>(uniform_dist(rng) * (upper - lower) + lower);
  }
}

T
tensor-tang 已提交
56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71
void vrelu_ref(const int n, const float* x, float* y) {
  for (int i = 0; i < n; ++i) {
    y[i] = x[i] > 0.f ? x[i] : 0.f;
  }
}

#if defined __AVX__ || defined __AVX2__
void vrelu_intri8(const int n, const float* x, float* y) {
  __m256 tmp = _mm256_loadu_ps(x);
  tmp = _mm256_max_ps(tmp, _mm256_setzero_ps());
  _mm256_storeu_ps(y, tmp);
}
#endif

TEST(JitKernel, vrelu) {
  namespace jit = paddle::operators::math::jitkernel;
T
tensor-tang 已提交
72
  for (int d : {3, 7, 8, 15, 16, 30, 256, 512}) {
T
tensor-tang 已提交
73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92
    std::vector<float> x(d);
    std::vector<float> zref(d), ztgt(d);
    RandomVec<float>(d, x.data(), -10.f, 1.f);
    const auto& ker =
        jit::KernelPool::Instance().template Get<jit::VReluKernel<float>>(d);
    const float* x_data = x.data();
    float* ztgt_data = ztgt.data();
    float* zref_data = zref.data();
    auto trefs = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
      vrelu_ref(d, x_data, zref_data);
    }
    auto trefe = GetCurrentUS();
#if defined __AVX__ || defined __AVX2__
    if (d == 8) {
      auto si0 = GetCurrentUS();
      for (int i = 0; i < repeat; ++i) {
        vrelu_intri8(d, x_data, zref_data);
      }
      auto si1 = GetCurrentUS();
93
      VLOG(30) << "Vec size 8 intr takes: " << (si1 - si0) / repeat;
T
tensor-tang 已提交
94 95 96 97
    }
#endif
    auto ttgts = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
T
tensor-tang 已提交
98
      ker->Compute(x_data, ztgt_data, d);
T
tensor-tang 已提交
99 100
    }
    auto ttgte = GetCurrentUS();
101 102 103
    VLOG(30) << "Vec size " << d
             << ": refer takes: " << (trefe - trefs) / repeat
             << " us, tgt takes: " << (ttgte - ttgts) / repeat;
T
tensor-tang 已提交
104 105 106 107 108 109
    for (int i = 0; i < d; ++i) {
      EXPECT_NEAR(ztgt_data[i], zref_data[i], 1e-3);
    }
  }
}

T
tensor-tang 已提交
110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134
void vaddbias_ref(const int n, const float a, const float* x, float* y) {
  for (int i = 0; i < n; ++i) {
    y[i] = x[i] + a;
  }
}

TEST(JitKernel, vaddbias) {
  namespace jit = paddle::operators::math::jitkernel;
  for (int d : {7, 8, 15, 16, 30, 64, 100, 128, 256}) {
    std::vector<float> x(d);
    std::vector<float> zref(d), ztgt(d);
    RandomVec<float>(d, x.data(), -2.f, 2.f);
    const auto& ker =
        jit::KernelPool::Instance().template Get<jit::VAddBiasKernel<float>>(d);
    const float a = 2.f;
    const float* x_data = x.data();
    float* ztgt_data = ztgt.data();
    float* zref_data = zref.data();
    auto trefs = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
      vaddbias_ref(d, a, x_data, zref_data);
    }
    auto trefe = GetCurrentUS();
    auto ttgts = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
T
tensor-tang 已提交
135
      ker->Compute(&a, x_data, ztgt_data, d);
T
tensor-tang 已提交
136 137 138
    }
    auto ttgte = GetCurrentUS();

139 140 141
    VLOG(30) << "Vec size " << d
             << ": refer takes: " << (trefe - trefs) / repeat
             << " us, tgt takes: " << (ttgte - ttgts) / repeat;
T
tensor-tang 已提交
142 143 144 145 146 147
    for (int i = 0; i < d; ++i) {
      EXPECT_NEAR(ztgt_data[i], zref_data[i], 1e-3);
    }
  }
}

T
tensor-tang 已提交
148 149 150 151 152 153 154 155 156 157 158 159 160 161
void vexp_ref(const int n, const float* x, float* y) {
  for (int i = 0; i < n; ++i) {
    y[i] = std::exp(x[i]);
  }
}

#ifdef PADDLE_WITH_MKLML
void vexp_mkl(const int n, const float* x, float* y) {
  paddle::platform::dynload::vsExp(n, x, y);
}
#endif

TEST(JitKernel, vexp) {
  namespace jit = paddle::operators::math::jitkernel;
T
tensor-tang 已提交
162
  for (int d : {1, 3, 4, 6, 7, 8, 12, 15, 16, 20, 30, 128, 256}) {
T
tensor-tang 已提交
163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186
    std::vector<float> x(d);
    std::vector<float> zref(d), ztgt(d);
    RandomVec<float>(d, x.data(), -2.f, 2.f);
    const auto& ker =
        jit::KernelPool::Instance().template Get<jit::VExpKernel<float>>(d);
    const float* x_data = x.data();
    float* ztgt_data = ztgt.data();
    float* zref_data = zref.data();
    auto trefs = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
      vexp_ref(d, x_data, zref_data);
    }
    auto trefe = GetCurrentUS();

#ifdef PADDLE_WITH_MKLML
    auto tmkls = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
      vexp_mkl(d, x_data, zref_data);
    }
    auto tmkle = GetCurrentUS();
#endif

    auto ttgts = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
T
tensor-tang 已提交
187
      // ker->Compute(x_data, ztgt_data);
T
tensor-tang 已提交
188
      ker->Compute(x_data, ztgt_data, d);
T
tensor-tang 已提交
189 190 191
    }
    auto ttgte = GetCurrentUS();

192 193
    VLOG(30) << "Vec size " << d
             << ": refer takes: " << (trefe - trefs) / repeat
T
tensor-tang 已提交
194
#ifdef PADDLE_WITH_MKLML
195
             << " us, mkl takes: " << (tmkle - tmkls) / repeat << " us, "
T
tensor-tang 已提交
196
#else
197
             << " us, "
T
tensor-tang 已提交
198
#endif
199
             << "tgt takes: " << (ttgte - ttgts) / repeat;
T
tensor-tang 已提交
200 201
    for (int i = 0; i < d; ++i) {
      EXPECT_NEAR(ztgt_data[i], zref_data[i], 1e-3);
202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228
    }
  }
}

inline float _sigmoid(float x) {
  const float min = SIGMOID_THRESHOLD_MIN;
  const float max = SIGMOID_THRESHOLD_MAX;
  float tmp = (x < min) ? min : ((x > max) ? max : x);
  return 1.f / (1.f + std::exp(-tmp));
}

void vsigmoid_ref(const int n, const float* x, float* y) {
  for (int i = 0; i < n; ++i) {
    y[i] = _sigmoid(x[i]);
  }
}

void vsigmoid_better(
    const std::shared_ptr<
        const paddle::operators::math::jitkernel::VExpKernel<float>>& vexp,
    const int n, const float* x, float* y) {
  const float min = SIGMOID_THRESHOLD_MIN;
  const float max = SIGMOID_THRESHOLD_MAX;
  for (int i = 0; i < n; ++i) {
    y[i] = (x[i] < min) ? min : ((x[i] > max) ? max : x[i]);
    y[i] = 0.f - y[i];
  }
T
tensor-tang 已提交
229
  vexp->Compute(y, y, n);
230 231 232 233 234 235 236
  for (int i = 0; i < n; ++i) {
    y[i] = 1.f / (1.f + y[i]);
  }
}

TEST(JitKernel, vsigmoid) {
  namespace jit = paddle::operators::math::jitkernel;
T
tensor-tang 已提交
237
  for (int d : {1, 3, 4, 6, 7, 8, 15, 16, 30, 32, 64, 100, 128, 256}) {
238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259
    std::vector<float> x(d);
    std::vector<float> zref(d), ztgt(d);
    RandomVec<float>(d, x.data(), -2.f, 2.f);
    const auto& ker =
        jit::KernelPool::Instance().template Get<jit::VSigmoidKernel<float>>(d);
    const auto& vexp =
        jit::KernelPool::Instance().template Get<jit::VExpKernel<float>>(d);
    const float* x_data = x.data();
    float* ztgt_data = ztgt.data();
    float* zref_data = zref.data();
    auto tmkls = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
      vsigmoid_better(vexp, d, x_data, zref_data);
    }
    auto tmkle = GetCurrentUS();
    auto trefs = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
      vsigmoid_ref(d, x_data, zref_data);
    }
    auto trefe = GetCurrentUS();
    auto ttgts = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
T
tensor-tang 已提交
260
      ker->Compute(x_data, ztgt_data, d);
261 262 263
    }
    auto ttgte = GetCurrentUS();

264 265 266 267
    VLOG(30) << "Vec size " << d
             << ": refer takes: " << (trefe - trefs) / repeat
             << " us, better(jit exp) takes: " << (tmkle - tmkls) / repeat
             << " us, tgt takes: " << (ttgte - ttgts) / repeat;
T
tensor-tang 已提交
268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291
    for (int i = 0; i < d; ++i) {
      EXPECT_NEAR(ztgt_data[i], zref_data[i], 1e-3);
    }
  }
}

inline float _tanh(float x) { return 2.f * _sigmoid(2.f * x) - 1.f; }

void vtanh_ref(const int n, const float* x, float* y) {
  for (int i = 0; i < n; ++i) {
    y[i] = _tanh(x[i]);
  }
}

void vtanh_better(
    const std::shared_ptr<
        const paddle::operators::math::jitkernel::VScalKernel<float>>& vscal,
    const std::shared_ptr<
        const paddle::operators::math::jitkernel::VSigmoidKernel<float>>&
        vsigmoid,
    const std::shared_ptr<
        const paddle::operators::math::jitkernel::VAddBiasKernel<float>>&
        vaddbias,
    const int n, const float* x, float* y) {
T
tensor-tang 已提交
292 293
  const float a = 2.f, b = -1.f;
  vscal->Compute(&a, x, y, n);
T
tensor-tang 已提交
294
  vsigmoid->Compute(y, y, n);
T
tensor-tang 已提交
295 296
  vscal->Compute(&a, y, y, n);
  vaddbias->Compute(&b, y, y, n);
T
tensor-tang 已提交
297 298 299 300
}

TEST(JitKernel, vtanh) {
  namespace jit = paddle::operators::math::jitkernel;
T
tensor-tang 已提交
301
  for (int d : {1, 2, 3, 4, 5, 6, 7, 8, 15, 16, 30, 32, 64, 100, 128, 256}) {
T
tensor-tang 已提交
302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327
    std::vector<float> x(d);
    std::vector<float> zref(d), ztgt(d);
    RandomVec<float>(d, x.data(), -2.f, 2.f);
    const auto& ker =
        jit::KernelPool::Instance().template Get<jit::VTanhKernel<float>>(d);
    const auto& vscal =
        jit::KernelPool::Instance().template Get<jit::VScalKernel<float>>(d);
    const auto& vsigmoid =
        jit::KernelPool::Instance().template Get<jit::VSigmoidKernel<float>>(d);
    const auto& vaddbias =
        jit::KernelPool::Instance().template Get<jit::VAddBiasKernel<float>>(d);
    const float* x_data = x.data();
    float* ztgt_data = ztgt.data();
    float* zref_data = zref.data();
    auto tmkls = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
      vtanh_better(vscal, vsigmoid, vaddbias, d, x_data, zref_data);
    }
    auto tmkle = GetCurrentUS();
    auto trefs = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
      vtanh_ref(d, x_data, zref_data);
    }
    auto trefe = GetCurrentUS();
    auto ttgts = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
T
tensor-tang 已提交
328
      ker->Compute(x_data, ztgt_data, d);
T
tensor-tang 已提交
329 330 331
    }
    auto ttgte = GetCurrentUS();

332 333 334 335
    VLOG(30) << "Vec size " << d
             << ": refer takes: " << (trefe - trefs) / repeat
             << " us, better(jit exp) takes: " << (tmkle - tmkls) / repeat
             << " us, tgt takes: " << (ttgte - ttgts) / repeat;
336 337
    for (int i = 0; i < d; ++i) {
      EXPECT_NEAR(ztgt_data[i], zref_data[i], 1e-3);
T
tensor-tang 已提交
338 339 340 341
    }
  }
}

T
tensor-tang 已提交
342 343 344 345 346 347 348 349 350
void lstm_ctht_ref(
    const std::shared_ptr<
        const paddle::operators::math::jitkernel::VSigmoidKernel<float>>&
        vsigmoid_3d,
    const std::shared_ptr<
        const paddle::operators::math::jitkernel::VTanhKernel<float>>& vtanh_d,
    const std::shared_ptr<
        const paddle::operators::math::jitkernel::VExpKernel<float>>& vexp_1,
    const int d, float* gates, const float* ct_1, float* ct, float* ht) {
T
tensor-tang 已提交
351 352
  vsigmoid_3d->Compute(gates + d, gates + d, 3 * d);
  vtanh_d->Compute(gates, gates, d);
T
tensor-tang 已提交
353 354 355 356 357 358 359 360 361
  const float *i = gates + d, *f = gates + d * 2, *o = gates + d * 3;
  const float min = SIGMOID_THRESHOLD_MIN;
  const float max = SIGMOID_THRESHOLD_MAX;
  for (int k = 0; k < d; ++k) {
    // C_t = C_t-1 * fgated + cand_gated * igated
    ct[k] = ct_1[k] * f[k] + gates[k] * i[k];
    // H_t = act_cell(C_t) * ogated
    float tmp = ct[k] * 2;
    tmp = 0.f - ((tmp < min) ? min : ((tmp > max) ? max : tmp));
T
tensor-tang 已提交
362
    vexp_1->Compute(&tmp, &tmp, 1);
T
tensor-tang 已提交
363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379
    tmp = 2.f / (1.f + tmp) - 1.f;
    ht[k] = tmp * o[k];
  }
}

void lstm_ctht_better(
    const std::shared_ptr<
        const paddle::operators::math::jitkernel::VSigmoidKernel<float>>&
        vsigmoid_3d,
    const std::shared_ptr<
        const paddle::operators::math::jitkernel::VTanhKernel<float>>& vtanh_d,
    const std::shared_ptr<
        const paddle::operators::math::jitkernel::VMulKernel<float>>& vmul_d,
    const std::shared_ptr<
        const paddle::operators::math::jitkernel::VAddKernel<float>>& vadd_d,
    const int d, float* gates, const float* ct_1, float* ct, float* ht) {
  int d2 = d * 2;
T
tensor-tang 已提交
380 381
  vsigmoid_3d->Compute(gates + d, gates + d, 3 * d);
  vtanh_d->Compute(gates, gates, d);
T
tensor-tang 已提交
382 383
  vmul_d->Compute(gates, gates + d, gates + d, d);
  vmul_d->Compute(ct_1, gates + d2, gates + d2, d);
T
tensor-tang 已提交
384
  vadd_d->Compute(gates + d, gates + d2, ct, d);
T
tensor-tang 已提交
385
  /* H_t = act_cell(C_t) * ogated */
T
tensor-tang 已提交
386
  vtanh_d->Compute(ct, gates + d2, d);
T
tensor-tang 已提交
387
  vmul_d->Compute(gates + d2, gates + d * 3, ht, d);
T
tensor-tang 已提交
388 389 390 391
}

TEST(JitKernel, lstm) {
  namespace jit = paddle::operators::math::jitkernel;
T
tensor-tang 已提交
392
  for (int d : {1, 2, 3, 4, 5, 6, 7, 8, 15, 16, 30, 32, 64, 100}) {
T
tensor-tang 已提交
393 394 395 396 397 398 399 400 401 402 403
    int d4 = d * 4;
    int d3 = d * 3;
    std::vector<float> x(d4), xref(d4);
    std::vector<float> ct_1(d), ct_tgt(d), ht_tgt(d);
    std::vector<float> ct_ref(d), ht_ref(d);
    RandomVec<float>(d4, x.data(), -2.f, 2.f);
    RandomVec<float>(d, ct_1.data(), -2.f, 2.f);
    memcpy(xref.data(), x.data(), sizeof(float) * d4);
    std::string act_gate = "sigmoid", act_cand = "tanh", act_cell = "tanh";
    const auto& ker =
        jit::KernelPool::Instance()
T
tensor-tang 已提交
404
            .template Get<jit::LSTMKernel<float>, const std::string&,
T
tensor-tang 已提交
405
                          const std::string&, const std::string&>(
T
tensor-tang 已提交
406
                act_gate, act_cand, act_cell, d, false);
T
tensor-tang 已提交
407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452
    // below kernels are used to compute refer
    const auto& vsigmoid_3d =
        jit::KernelPool::Instance().template Get<jit::VSigmoidKernel<float>>(
            d3);
    const auto& vtanh_d =
        jit::KernelPool::Instance().template Get<jit::VTanhKernel<float>>(d);
    const auto& vexp_1 =
        jit::KernelPool::Instance().template Get<jit::VExpKernel<float>>(1);
    const auto& vmul_d =
        jit::KernelPool::Instance().template Get<jit::VMulKernel<float>>(d);
    const auto& vadd_d =
        jit::KernelPool::Instance().template Get<jit::VAddKernel<float>>(d);

    float* x_data = x.data();
    float* xref_data = xref.data();
    const float* ct_1_data = ct_1.data();
    float* ct_tgt_data = ct_tgt.data();
    float* ht_tgt_data = ht_tgt.data();
    float* ct_ref_data = ct_ref.data();
    float* ht_ref_data = ht_ref.data();
    // compute once to check correctness
    lstm_ctht_ref(vsigmoid_3d, vtanh_d, vexp_1, d, xref_data, ct_1_data,
                  ct_ref_data, ht_ref_data);
    ker->ComputeCtHt(x_data, ct_1_data, ct_tgt_data, ht_tgt_data);
    for (int i = 0; i < d; ++i) {
      EXPECT_NEAR(ct_tgt_data[i], ct_ref_data[i], 1e-3);
      EXPECT_NEAR(ht_tgt_data[i], ht_ref_data[i], 1e-3);
    }

    auto tmkls = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
      lstm_ctht_better(vsigmoid_3d, vtanh_d, vmul_d, vadd_d, d, xref_data,
                       ct_1_data, ct_ref_data, ht_ref_data);
    }
    auto tmkle = GetCurrentUS();
    auto trefs = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
      lstm_ctht_ref(vsigmoid_3d, vtanh_d, vexp_1, d, xref_data, ct_1_data,
                    ct_ref_data, ht_ref_data);
    }
    auto trefe = GetCurrentUS();
    auto ttgts = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
      ker->ComputeCtHt(x_data, ct_1_data, ct_tgt_data, ht_tgt_data);
    }
    auto ttgte = GetCurrentUS();
453 454 455 456
    VLOG(30) << "Vec size " << d
             << ": refer takes: " << (trefe - trefs) / repeat
             << " us, better(jit) takes: " << (tmkle - tmkls) / repeat
             << " us, tgt takes: " << (ttgte - ttgts) / repeat;
T
tensor-tang 已提交
457 458 459
  }
}

T
tensor-tang 已提交
460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537
void vscal_ref(const int n, const float a, const float* x, float* y) {
  for (int i = 0; i < n; ++i) {
    y[i] = a * x[i];
  }
}
void vscal_inp_ref(const int n, const float a, float* x) {
  for (int i = 0; i < n; ++i) {
    x[i] = a * x[i];
  }
}
#if defined __AVX__ || defined __AVX2__
void vscal_intri8(const int n, const float a, const float* x, float* y) {
  __m256 tmp;
  __m256 scalar = _mm256_set1_ps(a);
  tmp = _mm256_loadu_ps(x);
  tmp = _mm256_mul_ps(tmp, scalar);
  _mm256_storeu_ps(y, tmp);
}
void vscal_inp_intri8(const int n, const float a, float* x) {
  __m256 tmp;
  __m256 scalar = _mm256_set1_ps(a);
  tmp = _mm256_loadu_ps(x);
  tmp = _mm256_mul_ps(tmp, scalar);
  _mm256_storeu_ps(x, tmp);
}
#endif

#ifdef PADDLE_WITH_MKLML
void vscal_inp_mkl(const int n, const float a, float* x) {
  paddle::platform::dynload::cblas_sscal(n, a, x, 1);
}
#endif

TEST(JitKernel, vscal) {
  namespace jit = paddle::operators::math::jitkernel;
  for (int d : {7, 8, 15, 16, 30, 256, 512}) {
    std::vector<float> x(d), y(d);
    std::vector<float> zref(d), ztgt(d);
    RandomVec<float>(d, x.data());
    std::memcpy(y.data(), x.data(), sizeof(float) * d);
    float a = 2.f;
    const auto& ker =
        jit::KernelPool::Instance().template Get<jit::VScalKernel<float>>(d);
    const float* x_data = x.data();
    float* y_data = y.data();
    float* ztgt_data = ztgt.data();
    float* zref_data = zref.data();
    auto trefs = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
      vscal_ref(d, a, x_data, zref_data);
    }
    auto trefe = GetCurrentUS();
    auto trefs1 = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
      vscal_inp_ref(d, a, y_data);
    }
    auto trefe1 = GetCurrentUS();

#ifdef PADDLE_WITH_MKLML
    auto tmkls = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
      vscal_inp_mkl(d, a, y_data);
    }
    auto tmkle = GetCurrentUS();
#endif

#if defined __AVX__ || defined __AVX2__
    if (d == 8) {
      auto si0 = GetCurrentUS();
      for (int i = 0; i < repeat; ++i) {
        vscal_intri8(d, a, x_data, zref_data);
      }
      auto si1 = GetCurrentUS();
      auto si2 = GetCurrentUS();
      for (int i = 0; i < repeat; ++i) {
        vscal_inp_intri8(d, a, y_data);
      }
      auto si3 = GetCurrentUS();
538 539
      VLOG(30) << "Vec size 8 intr takes: " << (si1 - si0) / repeat
               << " us, inplace: " << (si3 - si2) / repeat;
T
tensor-tang 已提交
540 541 542 543 544
    }
#endif

    auto ttgts = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
T
tensor-tang 已提交
545
      ker->Compute(&a, x_data, ztgt_data, d);
T
tensor-tang 已提交
546 547 548 549
    }
    auto ttgte = GetCurrentUS();
    auto ttgts1 = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
T
tensor-tang 已提交
550
      ker->Compute(&a, y_data, y_data, d);
T
tensor-tang 已提交
551 552
    }
    auto ttgte1 = GetCurrentUS();
553 554 555
    VLOG(30) << "Vec size " << d
             << ": refer takes: " << (trefe - trefs) / repeat
             << " us, inplace takes: " << (trefe1 - trefs1) / repeat
T
tensor-tang 已提交
556
#ifdef PADDLE_WITH_MKLML
557 558
             << " us, mkl inplace takes: " << (tmkle - tmkls) / repeat
             << " us, "
T
tensor-tang 已提交
559
#else
560
             << " us, "
T
tensor-tang 已提交
561
#endif
562 563
             << "tgt takes: " << (ttgte - ttgts) / repeat
             << "us, tgt inplace takes: " << (ttgte1 - ttgts1) / repeat;
T
tensor-tang 已提交
564 565 566 567 568
    for (int i = 0; i < d; ++i) {
      EXPECT_NEAR(ztgt_data[i], zref_data[i], 1e-3);
    }
  }
}
T
tensor-tang 已提交
569

T
tensor-tang 已提交
570 571 572 573 574 575
void vmul_ref(const int n, const float* x, const float* y, float* z) {
  for (int i = 0; i < n; ++i) {
    z[i] = x[i] * y[i];
  }
}

T
tensor-tang 已提交
576
#if defined __AVX__ || defined __AVX2__
T
tensor-tang 已提交
577
void vmul_intri8(const int n, const float* x, const float* y, float* z) {
T
tensor-tang 已提交
578 579 580 581 582 583 584
  __m256 tmpx, tmpy;
  tmpx = _mm256_loadu_ps(x);
  tmpy = _mm256_loadu_ps(y);
  tmpx = _mm256_mul_ps(tmpx, tmpy);
  _mm256_storeu_ps(z, tmpx);
}
#endif
T
tensor-tang 已提交
585

T
tensor-tang 已提交
586 587 588
#ifdef PADDLE_WITH_MKLML
void vmul_mkl(const int n, const float* x, const float* y, float* z) {
  paddle::platform::dynload::vsMul(n, x, y, z);
T
tensor-tang 已提交
589
}
T
tensor-tang 已提交
590
#endif
T
tensor-tang 已提交
591

T
tensor-tang 已提交
592 593
TEST(JitKernel, vmul) {
  namespace jit = paddle::operators::math::jitkernel;
594
  for (int d : {7, 8, 15, 16, 20, 30, 256, 512, 1000, 1024}) {
T
tensor-tang 已提交
595 596 597 598 599 600 601 602 603 604
    std::vector<float> x(d), y(d);
    std::vector<float> zref(d), ztgt(d);
    RandomVec<float>(d, x.data());
    RandomVec<float>(d, y.data());
    const auto& ker =
        jit::KernelPool::Instance().template Get<jit::VMulKernel<float>>(d);
    const float* x_data = x.data();
    const float* y_data = y.data();
    float* ztgt_data = ztgt.data();
    float* zref_data = zref.data();
T
tensor-tang 已提交
605
    auto trefs = GetCurrentUS();
T
tensor-tang 已提交
606
    for (int i = 0; i < repeat; ++i) {
T
tensor-tang 已提交
607
      vmul_ref(d, x_data, y_data, zref_data);
T
tensor-tang 已提交
608
    }
T
tensor-tang 已提交
609
    auto trefe = GetCurrentUS();
T
tensor-tang 已提交
610

T
tensor-tang 已提交
611 612
#ifdef PADDLE_WITH_MKLML
    auto tmkls = GetCurrentUS();
T
tensor-tang 已提交
613
    for (int i = 0; i < repeat; ++i) {
T
tensor-tang 已提交
614
      vmul_mkl(d, x_data, y_data, zref_data);
T
tensor-tang 已提交
615
    }
T
tensor-tang 已提交
616 617
    auto tmkle = GetCurrentUS();
#endif
T
tensor-tang 已提交
618

T
tensor-tang 已提交
619 620 621 622
#if defined __AVX__ || defined __AVX2__
    if (d == 8) {
      auto si0 = GetCurrentUS();
      for (int i = 0; i < repeat; ++i) {
T
tensor-tang 已提交
623
        vmul_intri8(d, x_data, y_data, zref_data);
T
tensor-tang 已提交
624 625
      }
      auto si1 = GetCurrentUS();
626
      VLOG(30) << "Vec size 8 intr takes: " << (si1 - si0) / repeat;
T
tensor-tang 已提交
627 628 629 630 631
    }
#endif

    auto ttgts = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
T
tensor-tang 已提交
632
      ker->Compute(x_data, y_data, ztgt_data, d);
T
tensor-tang 已提交
633 634 635
    }
    auto ttgte = GetCurrentUS();

636 637
    VLOG(30) << "Vec size " << d
             << ": refer takes: " << (trefe - trefs) / repeat
T
tensor-tang 已提交
638
#ifdef PADDLE_WITH_MKLML
639
             << " us, mkl takes: " << (tmkle - tmkls) / repeat << " us, "
T
tensor-tang 已提交
640
#else
641
             << " us, "
T
tensor-tang 已提交
642
#endif
643
             << "tgt takes: " << (ttgte - ttgts) / repeat;
T
tensor-tang 已提交
644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705
    for (int i = 0; i < d; ++i) {
      EXPECT_NEAR(ztgt_data[i], zref_data[i], 1e-3);
    }
  }
}

void vadd_ref(const int n, const float* x, const float* y, float* z) {
  for (int i = 0; i < n; ++i) {
    z[i] = x[i] + y[i];
  }
}

#if defined __AVX__ || defined __AVX2__
void vadd_intri8(const int n, const float* x, const float* y, float* z) {
  __m256 tmpx, tmpy;
  tmpx = _mm256_loadu_ps(x);
  tmpy = _mm256_loadu_ps(y);
  tmpx = _mm256_add_ps(tmpx, tmpy);
  _mm256_storeu_ps(z, tmpx);
}
#endif

#ifdef PADDLE_WITH_MKLML
void vadd_mkl(const int n, const float* x, const float* y, float* z) {
  paddle::platform::dynload::vsAdd(n, x, y, z);
}
#endif

TEST(JitKernel, vadd) {
  namespace jit = paddle::operators::math::jitkernel;
  for (int d : {7, 8, 15, 16, 30, 256, 512}) {
    std::vector<float> x(d), y(d);
    std::vector<float> zref(d), ztgt(d);
    RandomVec<float>(d, x.data());
    RandomVec<float>(d, y.data());
    const auto& ker =
        jit::KernelPool::Instance().template Get<jit::VAddKernel<float>>(d);
    const float* x_data = x.data();
    const float* y_data = y.data();
    float* ztgt_data = ztgt.data();
    float* zref_data = zref.data();
    auto trefs = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
      vadd_ref(d, x_data, y_data, zref_data);
    }
    auto trefe = GetCurrentUS();

#ifdef PADDLE_WITH_MKLML
    auto tmkls = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
      vadd_mkl(d, x_data, y_data, zref_data);
    }
    auto tmkle = GetCurrentUS();
#endif

#if defined __AVX__ || defined __AVX2__
    if (d == 8) {
      auto si0 = GetCurrentUS();
      for (int i = 0; i < repeat; ++i) {
        vadd_intri8(d, x_data, y_data, zref_data);
      }
      auto si1 = GetCurrentUS();
706
      VLOG(30) << "Vec size 8 intr takes: " << (si1 - si0) / repeat;
T
tensor-tang 已提交
707 708 709
    }
#endif

T
tensor-tang 已提交
710 711
    auto ttgts = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
T
tensor-tang 已提交
712
      ker->Compute(x_data, y_data, ztgt_data, d);
T
tensor-tang 已提交
713 714 715
    }
    auto ttgte = GetCurrentUS();

716 717
    VLOG(30) << "Vec size " << d
             << ": refer takes: " << (trefe - trefs) / repeat
T
tensor-tang 已提交
718
#ifdef PADDLE_WITH_MKLML
719
             << " us, mkl takes: " << (tmkle - tmkls) / repeat << " us, "
T
tensor-tang 已提交
720
#else
721
             << " us, "
T
tensor-tang 已提交
722
#endif
723
             << "tgt takes: " << (ttgte - ttgts) / repeat;
T
tensor-tang 已提交
724 725 726 727 728 729
    for (int i = 0; i < d; ++i) {
      EXPECT_NEAR(ztgt_data[i], zref_data[i], 1e-3);
    }
  }
}

T
tensor-tang 已提交
730 731 732 733 734 735 736 737 738 739 740
void vaddrelu_ref(const int n, const float* x, const float* y, float* z) {
  for (int i = 0; i < n; ++i) {
    z[i] = x[i] + y[i];
    z[i] = z[i] > 0 ? z[i] : 0;
  }
}
void vaddrelu_better(
    const std::shared_ptr<
        const paddle::operators::math::jitkernel::VAddKernel<float>>& vadd,
    const std::shared_ptr<
        const paddle::operators::math::jitkernel::VReluKernel<float>>& vrelu,
T
tensor-tang 已提交
741 742
    const float* x, const float* y, float* z, int d) {
  vadd->Compute(x, y, z, d);
T
tensor-tang 已提交
743
  vrelu->Compute(z, z, d);
T
tensor-tang 已提交
744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764
}

TEST(JitKernel, vaddrelu) {
  namespace jit = paddle::operators::math::jitkernel;
  for (int d : {7, 8, 15, 16, 30, 256, 512}) {
    std::vector<float> x(d), y(d);
    std::vector<float> zref(d), ztgt(d);
    RandomVec<float>(d, x.data());
    RandomVec<float>(d, y.data());
    const auto& ker =
        jit::KernelPool::Instance().template Get<jit::VAddReluKernel<float>>(d);
    const auto& vadd =
        jit::KernelPool::Instance().template Get<jit::VAddKernel<float>>(d);
    const auto& vrelu =
        jit::KernelPool::Instance().template Get<jit::VReluKernel<float>>(d);
    const float* x_data = x.data();
    const float* y_data = y.data();
    float* ztgt_data = ztgt.data();
    float* zref_data = zref.data();
    auto trefs = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
T
tensor-tang 已提交
765
      vaddrelu_ref(d, x_data, y_data, zref_data);
T
tensor-tang 已提交
766 767 768 769
    }
    auto trefe = GetCurrentUS();
    auto tmkls = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
T
tensor-tang 已提交
770
      vaddrelu_better(vadd, vrelu, x_data, y_data, zref_data, d);
T
tensor-tang 已提交
771 772 773 774
    }
    auto tmkle = GetCurrentUS();
    auto ttgts = GetCurrentUS();
    for (int i = 0; i < repeat; ++i) {
T
tensor-tang 已提交
775
      ker->Compute(x_data, y_data, ztgt_data, d);
T
tensor-tang 已提交
776 777
    }
    auto ttgte = GetCurrentUS();
778 779 780 781
    VLOG(30) << "Vec size " << d
             << ": refer takes: " << (trefe - trefs) / repeat
             << " us, better takes: " << (tmkle - tmkls) / repeat << " us, "
             << "tgt takes: " << (ttgte - ttgts) / repeat;
T
tensor-tang 已提交
782 783 784 785 786 787
    for (int i = 0; i < d; ++i) {
      EXPECT_NEAR(ztgt_data[i], zref_data[i], 1e-3);
    }
  }
}

T
tensor-tang 已提交
788 789 790 791
TEST(JitKernel, pool) {
  namespace jit = paddle::operators::math::jitkernel;
  const int frame_size = 4;
  std::string act_gate = "sigmoid", act_cand = "tanh", act_cell = "tanh";
T
tensor-tang 已提交
792
  const auto& plstm1 =
T
tensor-tang 已提交
793
      jit::KernelPool::Instance()
T
tensor-tang 已提交
794
          .template Get<jit::LSTMKernel<float>, const std::string&,
T
tensor-tang 已提交
795
                        const std::string&, const std::string&>(
T
tensor-tang 已提交
796
              act_gate, act_cand, act_cell, frame_size, false);
T
tensor-tang 已提交
797
  const auto& plstm2 =
T
tensor-tang 已提交
798
      jit::KernelPool::Instance()
T
tensor-tang 已提交
799
          .template Get<jit::LSTMKernel<float>, const std::string&,
T
tensor-tang 已提交
800
                        const std::string&, const std::string&>(
T
tensor-tang 已提交
801 802 803 804 805 806 807
              act_gate, act_cand, act_cell, frame_size, false);
  const auto& peephole =
      jit::KernelPool::Instance()
          .template Get<jit::LSTMKernel<float>, const std::string&,
                        const std::string&, const std::string&>(
              act_gate, act_cand, act_cell, frame_size, true);
  EXPECT_TRUE(plstm1 != peephole);
T
tensor-tang 已提交
808

T
tensor-tang 已提交
809
  const auto& pvmul_f =
T
tensor-tang 已提交
810
      jit::KernelPool::Instance().template Get<jit::VMulKernel<float>>(4);
T
tensor-tang 已提交
811 812
  EXPECT_TRUE(std::dynamic_pointer_cast<const jit::Kernel>(plstm2) !=
              std::dynamic_pointer_cast<const jit::Kernel>(pvmul_f));
T
tensor-tang 已提交
813

T
tensor-tang 已提交
814
  const auto& pvmul_d =
T
tensor-tang 已提交
815
      jit::KernelPool::Instance().template Get<jit::VMulKernel<double>>(4);
T
tensor-tang 已提交
816 817
  EXPECT_TRUE(std::dynamic_pointer_cast<const jit::Kernel>(pvmul_f) !=
              std::dynamic_pointer_cast<const jit::Kernel>(pvmul_d));
T
tensor-tang 已提交
818

T
tensor-tang 已提交
819
  const auto& pvmul_from_key = jit::KernelPool::Instance().Get("vmulfjit4");
T
tensor-tang 已提交
820 821 822 823 824
#if defined(__APPLE__) || defined(__OSX__) || defined(_WIN32)
  EXPECT_EQ(pvmul_from_key, nullptr);
#else
  EXPECT_EQ(pvmul_from_key, pvmul_f);
#endif
T
tensor-tang 已提交
825
  const auto& pvmul_from_key2 = jit::KernelPool::Instance().Get("vmulfjit");
T
tensor-tang 已提交
826
  EXPECT_TRUE(pvmul_from_key2 == nullptr);
T
tensor-tang 已提交
827
}