diff options
author | Jianyu Huang <jianyuhuang@fb.com> | 2019-01-31 02:06:15 +0300 |
---|---|---|
committer | Facebook Github Bot <facebook-github-bot@users.noreply.github.com> | 2019-01-31 02:17:18 +0300 |
commit | 03a8fa506ee0822e3ab045f8e3df902abcfd2e74 (patch) | |
tree | 95a5b96dfad4707fd006a9e7e00fad9becd63c6d /bench | |
parent | 79333308f5e2fc242727879dcd3de3536b6ffc39 (diff) |
Add threading for FBGEMM FP16
Summary: Add threading support for FBGEMM FP16 routines.
Reviewed By: dskhudia, jacobkahn
Differential Revision: D13792341
fbshipit-source-id: eb31a11340ac9fd0ee9b4f570d161e7c7e6a7602
Diffstat (limited to 'bench')
-rw-r--r-- | bench/FP16Benchmark.cc | 39 |
1 files changed, 35 insertions, 4 deletions
diff --git a/bench/FP16Benchmark.cc b/bench/FP16Benchmark.cc index efb043d..f3a2e2d 100644 --- a/bench/FP16Benchmark.cc +++ b/bench/FP16Benchmark.cc @@ -122,8 +122,22 @@ void performance_test() { C_ref.data(), n); #endif - cblas_gemm_compute( - matrix_op_t::NoTranspose, m, A.data(), Bp, beta, C_fb.data()); +#ifdef _OPENMP +#pragma omp parallel +#endif + { + int num_threads = fbgemm_get_num_threads(); + int tid = fbgemm_get_thread_num(); + cblas_gemm_compute( + matrix_op_t::NoTranspose, + m, + A.data(), + Bp, + beta, + C_fb.data(), + tid, + num_threads); + } #if defined(USE_MKL) || defined(USE_BLAS) // Compare results @@ -201,8 +215,25 @@ void performance_test() { } t_begin = chrono::system_clock::now(); - cblas_gemm_compute( - matrix_op_t::NoTranspose, m, A.data(), Bp, beta, C_fb.data()); + +#ifdef _OPENMP +#pragma omp parallel +#endif + { + int num_threads = fbgemm_get_num_threads(); + int tid = fbgemm_get_thread_num(); + + cblas_gemm_compute( + matrix_op_t::NoTranspose, + m, + A.data(), + Bp, + beta, + C_fb.data(), + tid, + num_threads); + } + t_end = chrono::system_clock::now(); if (it >= 0) { |