You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by ak...@apache.org on 2021/09/02 07:28:57 UTC
[incubator-mxnet] branch v1.x updated: [v1.x] Enabling BRGEMM
FullyConnected based on shapes (#20533)
This is an automated email from the ASF dual-hosted git repository.
akarbown pushed a commit to branch v1.x
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/v1.x by this push:
new 59e9b94 [v1.x] Enabling BRGEMM FullyConnected based on shapes (#20533)
59e9b94 is described below
commit 59e9b940413e18457f44c4951bf27236cbed4e8d
Author: bgawrych <ba...@intel.com>
AuthorDate: Thu Sep 2 09:27:00 2021 +0200
[v1.x] Enabling BRGEMM FullyConnected based on shapes (#20533)
* Enable brgemm based on input info
* fix sanity
* Review fixes
* Change function name
* Fix typo
* Align variable assignments
* Fix review
* use const reference
---
docs/static_site/src/pages/api/faq/env_var.md | 6 +++---
src/operator/nn/mkldnn/mkldnn_base-inl.h | 19 +++++++++++++++----
src/operator/nn/mkldnn/mkldnn_fully_connected.cc | 11 ++++++-----
3 files changed, 24 insertions(+), 12 deletions(-)
diff --git a/docs/static_site/src/pages/api/faq/env_var.md b/docs/static_site/src/pages/api/faq/env_var.md
index 6336b3d..b9e0c95 100644
--- a/docs/static_site/src/pages/api/faq/env_var.md
+++ b/docs/static_site/src/pages/api/faq/env_var.md
@@ -298,9 +298,9 @@ If ctypes is used, it must be `mxnet._ctypes.ndarray.NDArrayBase`.
- Values: Int ```(default=-1)```
- Flag to set num of elements that MKLDNN cache can hold. Default is -1 which means cache size is unbounded. Should only be set if your model has variable input shapes, as cache size may grow unbounded. The number represents the number of items in the cache and is proportional to the number of layers that use MKLDNN and different input shape.
-* MXNET_MKLDNN_DISABLE_BRGEMM_FC
- - Values: 0, 1 ```(default=1)```
- - Flag which disables BRGEMM kernels in FullyConnected executed with MKLDNN support - Should only be set to 0 if your model has constant input shapes or FullyConnected is calculated with large tensors. Supported on machines with AVX512-VNNI.
+* MXNET_MKLDNN_FORCE_FC_AB_FORMAT
+ - Values: 0, 1 ```(default=0)```
+ - If set to true, FullyConnected will use only AB format for weights, thus MXNet won't use BRGEMM implementation of FC on machines with AVX512-VNNI support which requires special weights format.
* MXNET_ENFORCE_DETERMINISM
- Values: 0(false) or 1(true) ```(default=0)```
diff --git a/src/operator/nn/mkldnn/mkldnn_base-inl.h b/src/operator/nn/mkldnn/mkldnn_base-inl.h
index 515d2f4..e206e9e 100644
--- a/src/operator/nn/mkldnn/mkldnn_base-inl.h
+++ b/src/operator/nn/mkldnn/mkldnn_base-inl.h
@@ -286,7 +286,16 @@ inline static mkldnn::memory::desc GetMemDesc(const NDArray& arr, int dtype = -1
return mkldnn::memory::desc{dims, get_mkldnn_type(dtype), mkldnn::memory::format_tag::any};
}
-inline static mkldnn::memory::desc GetFCWeightDesc(const NDArray& arr, int dtype = -1) {
+inline static bool ChooseBRGEMMImpl(const mkldnn::memory::dims& weight_dims, size_t batch_size) {
+ // Conditions based on measurement results done on CLX8280
+ // https://github.com/apache/incubator-mxnet/pull/20533
+ return weight_dims[0] >= 1024 && weight_dims[1] >= 1024 && batch_size >= 16384 &&
+ weight_dims[0] % 64 == 0 && weight_dims[1] % 64 == 0;
+}
+
+inline static mkldnn::memory::desc GetFCWeightDesc(const NDArray& arr,
+ size_t batch_size,
+ int dtype = -1) {
int ndim = arr.shape().ndim();
mkldnn::memory::dims dims(ndim);
dtype = (dtype == -1) ? arr.dtype() : dtype;
@@ -294,9 +303,11 @@ inline static mkldnn::memory::desc GetFCWeightDesc(const NDArray& arr, int dtype
dims[i] = arr.shape()[i];
auto format = mkldnn::memory::format_tag::any;
// for batch 256 alexnet benchmark test
- const bool brgemm_disabled = dmlc::GetEnv("MXNET_MKLDNN_DISABLE_BRGEMM_FC", true);
- if (dims.size() == 2 && brgemm_disabled) {
- format = mkldnn::memory::format_tag::ab;
+ const bool force_fc_ab_format = dmlc::GetEnv("MXNET_MKLDNN_FORCE_FC_AB_FORMAT", false);
+ if (dims.size() == 2) {
+ if (force_fc_ab_format || !ChooseBRGEMMImpl(dims, batch_size)) {
+ format = mkldnn::memory::format_tag::ab;
+ }
}
return mkldnn::memory::desc{dims, get_mkldnn_type(dtype), format};
diff --git a/src/operator/nn/mkldnn/mkldnn_fully_connected.cc b/src/operator/nn/mkldnn/mkldnn_fully_connected.cc
index 4e0aa5a..ccd1287 100644
--- a/src/operator/nn/mkldnn/mkldnn_fully_connected.cc
+++ b/src/operator/nn/mkldnn/mkldnn_fully_connected.cc
@@ -41,10 +41,11 @@ mkldnn::inner_product_forward::primitive_desc GetFCFwdImpl(const MKLDNNFCFullPar
const NDArray& weight,
const NDArray* bias,
const mkldnn::memory::desc& out_md) {
- auto data_md = GetMemDesc(data);
- auto weight_md = full_param.mkldnn_param.quantized ? GetFCWeightDesc(weight, mshadow::kInt8)
- : GetFCWeightDesc(weight);
auto engine = CpuEngine::Get()->get_engine();
+ auto data_md = GetMemDesc(data);
+ auto weight_md = full_param.mkldnn_param.quantized
+ ? GetFCWeightDesc(weight, data.shape()[0], mshadow::kInt8)
+ : GetFCWeightDesc(weight, data.shape()[0]);
auto propagation =
is_train ? mkldnn::prop_kind::forward_training : mkldnn::prop_kind::forward_scoring;
@@ -107,7 +108,7 @@ inline static mkldnn::inner_product_backward_data::primitive_desc GetFCBwdData(
const NDArray& output,
mkldnn::inner_product_forward::primitive_desc fwd_pd) {
auto data_md = GetMemDesc(data);
- auto weight_md = GetFCWeightDesc(weight);
+ auto weight_md = GetFCWeightDesc(weight, data.shape()[0]);
auto out_md = GetMemDesc(output);
auto engine = CpuEngine::Get()->get_engine();
mkldnn::inner_product_backward_data::desc desc(data_md, weight_md, out_md);
@@ -121,7 +122,7 @@ inline static mkldnn::inner_product_backward_weights::primitive_desc GetFCBwdWei
const NDArray& output,
mkldnn::inner_product_forward::primitive_desc fwd_pd) {
auto data_md = GetMemDesc(data);
- auto weight_md = GetFCWeightDesc(weight);
+ auto weight_md = GetFCWeightDesc(weight, data.shape()[0]);
auto out_md = GetMemDesc(output);
auto engine = CpuEngine::Get()->get_engine();
if (bias) {