Integrating the MKL VML functions to MXNET to speed-up the (element-wised) mathematic computation

[juliusshufan]

Description

Intel MKL provides a wide range of the generic vectored mathematic (VML) functions befitting from AVX512 instructions, via the integration of the VML functions, some of the element-wised operation is expected to be speedup.

Currently, the generic mathematic computations are implemented by a series of mshadow OPs, which encapsulate the functions provided by standard math library, and the vectorization of the input is implemented by dense/sparse tensor, and computations are paralized by OpenMP. Specially,
Element-wise OP taking one or two inputs, with “write-to”/”write-inplace” types can be supported by MKL VML functions.

@TaoLv @pengzhao-intel

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [MXNET-$JIRA_ID], where $JIRA_ID refers to the relevant JIRA issue created (except PRs with tiny changes)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage:
• Unit tests are added for small changes to verify correctness (e.g. adding a new operator)
• Nightly tests are added for complicated/long-running ones (e.g. changing distributed kvstore)
• Build tests will be added for build configuration changes (e.g. adding a new build option with NCCL)
• Code is well-documented:
• For user-facing API changes, API doc string has been updated.
• For new C++ functions in header files, their functionalities and arguments are documented.
• For new examples, README.md is added to explain the what the example does, the source of the dataset, expected performance on test set and reference to the original paper if applicable
• Check the API doc at http://mxnet-ci-doc.s3-accelerate.dualstack.amazonaws.com/PR-$PR_ID/$BUILD_ID/index.html
• To the my best knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

The unit tests are provided with the changes.

[anirudhacharya]

@mxnet-label-bot add [pr-awaiting-review]

[pengzhao-intel]

@eric-haibin-lin @szha please help take a review.

[pengzhao-intel]

Does all of these functions will be mapped automatically when MKL is enabled?

[TaoLv]

No. We just put all the VML functions here. We think these functions can be leveraged by MXNet in the future. But currently it need to change the registration of each operator to use these functions. In this PR we only optimized some operators which are used in BERT. We propose to optimize others when we face performance problems on them.

[pengzhao-intel]

Thanks for the explanation. We can add it back when we use it; otherwise, it is a little confusion for other developers.

[eric-haibin-lin]

Why do you override the FComputeEx attribute?

[TaoLv]

@eric-haibin-lin Thanks for review. Not sure how sparse is handled in the original FComputeEx. Previously I thought sparse can also benefit from VML if its values are stored in a dense way. But we don't much data to prove that.

[eric-haibin-lin]

Thanks for the explanation. Yeah it should benefit

[TaoLv]

I just reverted the change for FComputeEx as we don't have much data for that yet. Will revisit this part once we meet any performance issue for sparse.

[eric-haibin-lin]

I actually I'd prefer not reverting it. I don't see any reason why it won't help. Let's undo the revert?

[TaoLv]

Can you provide a simple benchmark for a sparse unary operator? We can take a quick try. Thanks! cc @juliusshufan

[juliusshufan]

@eric-haibin-lin @TaoLv Friendly ping, May I know your decison on the sparse part?

[eric-haibin-lin]

Does this PR also want to enable optimization for layernorm ?

[TaoLv]

Yes. I'm working on enabling it and trying to understand the optimization and workflow from @sxjscience 's PR.

comments addressed

[pengzhao-intel]

@eric-haibin-lin please help to review again :)

[pengzhao-intel]

The minor comment is added.

LGTM

[pengzhao-intel]

Any chance to fusion some of these operations to reduce the memory bandwidth?

[eric-haibin-lin]

How much faster is this version compared to the mshadow one?

[sxjscience]

After reading the code, I think the current implementation, which relies on the vectorized operations, should be fast at scaling and shifting the data (data * gamma & data + beta). One possible improvement is to use the Welford's online algorithm (https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance) to calculate the mean/variance in one pass, the code will look like this:

template <typename DType>
MSHADOW_XINLINE static void mean_var_(index_t n, DType *in, DType *mean, DType* variance) {
  DType sigma2 = 0;
  DType mean_v = 0;
  DType old_mean_v = 0;
  for (index_t i = 0; i < n; i++) {
    DType x = in[i];
    old_mean_v = mean_v;
    mean_v += (x - old_mean_v) / (i + 1);
    sigma2 += (x - old_mean_v) * (x - mean_v);
  }    
  mean[0] = mean_v;
  variance[0] = sigma2 / n;
}


template <typename DType>
MSHADOW_XINLINE static void LayerNormLastDim(index_t m,
                                             index_t n,
                                             DType *a,
                                             DType *b,
                                             DType *ws,
                                             DType *gamma,
                                             DType *beta,
                                             DType *mean,
                                             DType *var,
                                             DType eps) {
  auto nthreads = engine::OpenMP::Get()->GetRecommendedOMPThreadCount();
#pragma omp parallel for num_threads(nthreads)
  for (index_t i = 0; i < m; i++) {
    DType ele_mean, ele_var;
    DType* in_offset = a + i * n;
    DType* out_offset = b + i * n;
    mean_var_(n, in_offset, &ele_mean, &ele_var);
    sub_(n, in_offset, ele_mean, out_offset);
    ele_var = math::sqrt(ele_var + eps);
    mul::Vectorize(n, out_offset, gamma, out_offset);
    div_(n, out_offset, ele_var, out_offset);
    add::Vectorize(n, out_offset, beta, out_offset);
    mean[i] = ele_mean;
    var[i] = ele_var;
  }
}

[TaoLv]

@pengzhao-intel @sxjscience loops are fused in the latest commit. I also removed the required workspace but that means we can not leverage VML functions and need rely on compiler for vectorization.

[TaoLv]

@sxjscience Can you help to review? Here is a optimization for CPU LayerNorm.

[TaoLv]

LayerNorm performance is measured on my skl machine. Shapes are from BERT base and large model respectively. The speedup from this PR is around 3x~10x. @eric-haibin-lin @sxjscience @pengzhao-intel

# mxnet-mkl==1.4.1
layernorm (1L, 128L, 768L): 0.23437 ms
layernorm (8L, 128L, 768L): 1.39641 ms
layernorm (32L, 128L, 768L): 5.18604 ms
layernorm (1L, 128L, 1024L): 0.35661 ms
layernorm (8L, 128L, 1024L): 1.80795 ms
layernorm (32L, 128L, 1024L): 6.76601 ms

# this PR built with USE_BLAS=mkl
layernorm (1, 128, 768): 0.07230 ms
layernorm (8, 128, 768): 0.21550 ms
layernorm (32, 128, 768): 0.51188 ms
layernorm (1, 128, 1024): 0.08863 ms
layernorm (8, 128, 1024): 0.25120 ms
layernorm (32, 128, 1024): 0.63479 ms

[pengzhao-intel]

@sxjscience @eric-haibin-lin any other comments? If no, I will merge this PR soon for the release 1.5.