Onednn: Slowdown of FP32->BF16 Reorder for Specific Size

[jiecaoyu]

Summary

I am using reorder to convert 32-bit floating-point (FP32) numbers into BF16. However, I found that, for some specific sizes, the reorder primitive suffers from a throughput degradation.

Version

oneDNN version: v1.2

Steps to reproduce

The test code (reorder_slowdown.cpp) is attached at the end. I tested three sizes:

• 140 * 3850 = 539000
• 140 * 3851 = 539140
• 140 * 3852 = 539280

The throughput results I can get are:

Reorder Throughput with Size 539000   1.8 GOPS
Reorder Throughput with Size 539140   0.5 GOPS
Reorder Throughput with Size 539280   1.9 GOPS

As you can see, the throughput for (140 * 3851 = 539140) is dramatically lower than that for other sizes.

Test Code

#include <iostream>
#include <random>
#include <vector>
#include <chrono>
#include <iomanip>

#include "mkldnn.hpp"

void init_vector(std::vector<float> &v) {
  std::mt19937 gen;
  std::uniform_real_distribution<float> u(-1, 1);
  for (auto &e : v)
    e = u(gen);
}

int main(int argc, char** argv) {
  // Tested shape
  const int M = 256;
  const int N = 256;
  std::vector<int> test_sizes {140 * 3850, 140 * 3851, 140 * 3852};

  const int NWAMRUP = 200;
  const int NITER = 1000;

  constexpr mkldnn::memory::data_type bf16 = mkldnn::memory::data_type::bf16;
  constexpr mkldnn::memory::data_type f32 = mkldnn::memory::data_type::f32;

  mkldnn::engine eng(mkldnn::engine::kind::cpu, 0);
  mkldnn::stream s(eng);

  // Cache flush
  std::vector<float> LLC;
  const int LLC_size = 40 * 1024 * 1024 / 4;
  LLC.resize(LLC_size, 1.0f);
  mkldnn::memory::desc LLC_md({LLC_size}, f32, mkldnn::memory::format_tag::a);
  mkldnn::eltwise_forward::desc llc_flush_d(mkldnn::prop_kind::forward_training,
      mkldnn::algorithm::eltwise_linear, LLC_md, 0.f, 1.f);
  mkldnn::eltwise_forward::primitive_desc llc_flush_pd(llc_flush_d, eng);
  mkldnn::eltwise_forward llc_flush_p(llc_flush_pd);
  mkldnn::memory LLC_m(LLC_md, eng, (void*)LLC.data());
  for (int size: test_sizes) {
    std::vector<float> X(size);
    init_vector(X);
    mkldnn::memory::desc X_md({size}, f32, mkldnn::memory::format_tag::a);
    mkldnn::memory::desc X_bf16_md({size}, bf16, mkldnn::memory::format_tag::a);
    mkldnn::memory X_m(X_md, eng, (void*)X.data());
    mkldnn::memory X_bf16_m(X_bf16_md, eng);

    mkldnn::reorder reorder_p(X_m, X_bf16_m);

    double ttot = 0.0;

    std::chrono::time_point<std::chrono::high_resolution_clock> start, end;
    for (int i = 0; i < (NWAMRUP + NITER); ++i) {
      llc_flush_p.execute(
          s,
          {{DNNL_ARG_SRC, LLC_m},
          {DNNL_ARG_DST, LLC_m}});
      s.wait();
      start = std::chrono::high_resolution_clock::now();
      reorder_p.execute(s, X_m, X_bf16_m);
      s.wait();
      end = std::chrono::high_resolution_clock::now();
      auto dur = std::chrono::duration_cast<std::chrono::nanoseconds>(end - start);
      if (i >= NWAMRUP) ttot += dur.count();
    }

    ttot /= NITER;

    std::cout << std::setw(22) << "Reorder Throughput with Size "
      << std::setw(4) << size << " "
      << std::fixed << std::setw(5) << std::setprecision(1)
      << size/ttot << " GOPS" << std::endl;
  }
  return 0;
}

[alokbakshi]

Hi @jiecaoyu

Thank you for the finding! Just by looking at the jit dump I could see lot SSE instructions for the particular test shapes that you've run. Perhaps that was (one of) the cause behind difference in performance that you observed. I will get back to you after preparing patches and running performance experiments on my end.

[alokbakshi]

Hi @jiecaoyu

I did get rid of the SSE instructions to avoid potential SSE-AVX transition penalty but unfortunately that did not have much impact on overall performance.

Based upon the low GOPS values you observed I assume that you have run the problem instance on a single thread/core. Running your example on 24 cores/socket Cooper Lake machine with commands Invoked:

# Use all 24 threads in first socket
KMP_AFFINITY=compact,granularity=core numactl --physcpubind=0-24 --membind=0 ./f32_bf16_reorder
# Use only single thread
KMP_AFFINITY=compact,granularity=core OMP_NUM_THREADS=1 numactl -m 0 -N 0  ./f32_bf16_reorder

the outcome I got is summarized below:

| Size | GOPS (24 Threads) | GOPS (1 Thread) |
|---------------|-----------------------|---------------------|
| 539000 | 38.6 | 3.4 |
| 539140 | 38.3 | 1.0 |
| 539280 | 37.7 | 3.7 |

Note: The results varied from run to run by not more than 6 %

So as you mentioned, for middle shape the tests verify the throughput degradation for single threaded instance.

Getting back to the particular shapes you have used

• 140 * 3850 = 539000
• 140 * 3851 = 539140
• 140 * 3852 = 539280

we see that 3851 is a prime number and that I think is the reason behind low performance. Diving deeper into the reorder code we can see that it breaks apart the problem into driver and jit code. For sake, of simplicity let us take the problem size of 140 * 3851 = 539140 then the bird-eye view of code would be

 # each thread will deal with bunch of contiguous indices
 for idx in range(0, 3851):
     inp = inp_data[idx * 140  : (idx + 1) * 140 ]
     out = out_data[idx * 140  : (idx + 1) * 140 ]

     # reorder 140 many of data elements by using optimized assembly code
     jit_ker(inp, out)

Note:

• If you want, you can un-comment #define TR_DEBUG in jit_uni_reorder.cpp file and rebuild the library to see the relevant values of these blocking parameters.
• In general, a problem of size N is decomposed into N = N_j X N_d = N_j X N_{d1} X N_{d2} ... X N_{dk} i.e. driver work is further blocked (and sometimes rearranged) to improve cache usage but that does not happen for prime N_d = 3851.

Via inspection, the jit kernel work per iteration can be found to be unusually low in problematic shape as seen below

| Size | Jit Kernel Work | Driver Work |
|---------------|-----------------------|----------------|
| 539000 | 1400 | 385 |
| 539140 | 140 | 3851 |
| 539280 | 1284 | 420 |

Therefore, for the second shape the driver work 539140 / 140 = 3851 is comparatively large resulting in in-efficient blocking for single-threaded case.

What is the work around? I think to fully optimize we need to break apart the problem additively (apart from multiplicatively) as well. But the resulting code-complexity has to be justified by the practical application. To keep the codebase manageable oftentimes the jit kernels are only optimized for the practically occurring shapes.

*edit: Fixed typos and added a clarifying note.

[Jianhui-Li]

Nice analysis. The overhead of invoking jit kernel might be one good reason. I also wonder why when you run this workload using 24 cores, it gets 38.3 times faster comparing to single core - it is a super-linear speedup.
539140 | 38.3 | 1.0

[alokbakshi]

Hi @Jianhui-Li and @jiecaoyu

As suggested, overhead of invoking jit kernel seems to be one of the reason behind this performance gap.

For OMP_NUM_THREADS=1 artificially making jit kernel deal with whole input in one go [i.e. there is no looping around the jit call] significantly improves the performance for 140 x 3851 problem size. For example, below is the sample run

Size | GOPS (1 Thread)
-- | -- |
539000 | 3.4
539140 | 3.7
539280 | 3.6

where the reorder code was made to do so by tweaking the constants (say 16 * nthr ---> (nthr > 1 ? 16 : 1) * nthr) in the line below

// sz_drv_min: minimal size for the parallel driver, required for good parallelization
size_t sz_drv_min = min(16 * nthr, div_up(sz_total, 1024));

Regarding the super-linear speed up: my guess is that -- owing to large problem size -- for single thread the cache eviction plays a role in performance degradation [have to look more carefully though to make definite statement]

I reran experiments to see how for problem with sizes 140 * P where P is a large prime number the speedup varies with respect to the problem size

Size | Time (1 Thread) | Time (24 Threads) | Speedup|
-- | -- | -- | -- |
140 x 1061 | 0.09 | 0.0040 | 22.5
140 x 2129 | 0.19 | 0.0049 | 38.7
140 x 3851 | 0.34 | 0.0068 | 50
140 x 5441 | 0.48 | 0.0081 | 59
140 x 6737 | 0.60 | 0.0095 | 63

and the results seem to verify the role of cache in super-linear speedup.

Note: just for reverification, the timings in the table above were measured via benchdnn tool and though it maybe less (or probably more) precise in gauging individual timing but still I think the same trend should hold up elsewhere.

All in all, I think we have to run comprehensive tests to find the optimal value of threading constants so that the changes do not lead to regression elsewhere. And with that the performance for the problematic shape should be at par with rest of the others.

I'd test it on my side and will let you know when the patch is ready.

PS: Thanks a lot for sharing your insight and results!

[alokbakshi]

*Addendum:
On my system (with L1 cache = 32 K and L2 cache of 1024 K) running the below command for shape 140 x 6737 = 943180 gives

for T in 0 2 5 8 11 14 17 20 23; do
KMP_AFFINITY=compact,granularity=core numactl --physcpubind=0-${T} --membind=0  ./benchdnn --reorder --sdt=bf16 --ddt=f32 --stag=x --dtag=x --mode=P 943180
done 

Threads | Time
-- | --
1 | 0.594
3 | 0.2048
6 | 0.154
9 | 0.020

Note: Rest of the timings seem to vary a lot from run to run and so not included above.

Furthermore, the total size of problem is 943180 x ( sizeof(bfloat16_t) + sizeof(float) ) = 5526 K i.e. around 6 times L2 cache size.

So if we use 9 cores then for each core the data can comfortably fit into the L2 cache giving the performance boost. And that's what we see in the reading above where the performance simply goes up by 10 times via increasing thread count from 3 to 9.

[jiecaoyu]

Hi @alokbakshi , thanks for your detailed analysis! I wonder whether there will be any plan to make the change of 16 * nthr ---> (nthr > 1 ? 16 : 1) * nthr?

Another advantage of making this change is that we can perform FP32->BF16 in-place. Here, by "in-place", I am not referring to the element-wise in-place operation. Instead, if the FP32 tensor will not be used for future computation, I am using the first half of the memory space for the FP32 tensor to store the corresponding BF16 tensor. This trick can help further reduce the memory footprint and help a lot in some cases.

[alokbakshi]

Hi @jiecaoyu
After finishing performance tests on my end (so as to ensure that the heuristic does not bring performance regression elsewhere) and review thereof the changes will be promoted to the master branch. I'd update you here as soon as it gets done.

Regarding fp32 -> bf16 memory-bound reorder, reusing memory sounds like a great idea!

[alokbakshi]

Hi @jiecaoyu
Just wanted to update that the change has been promoted to the master branch (commit 4e5502c0b678c3d5e7cf97d8d9535b33a20ee54b)