-
Notifications
You must be signed in to change notification settings - Fork 6.8k
MKLDNN can be turned off with env var #12058
Conversation
|
0d1829f
to
03bd210
Compare
if (dev_mask == mshadow::cpu::kDevMask | ||
if (!MKLDNNEnvSet()) { | ||
*dispatch_mode = DispatchMode::kFComputeFallback; | ||
} else if (dev_mask == mshadow::cpu::kDevMask |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Please also align the 2 lines below with this one.
Maybe the env variable should start with Another question, what's the behavior of enabling this variable when MKL-DNN is not compiled into MXNet? Should we give a proper warning or tip for that? |
@@ -137,6 +137,10 @@ static inline bool SupportMKLDNN(const NDArray &input) { | |||
&& SupportStorageMKLDNN(input.storage_type()); | |||
} | |||
|
|||
static inline bool MKLDNNEnvSet() { | |||
return dmlc::GetEnv("USE_MKLDNN", true); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Seems like you're setting the default to be true
, would you please also add corresponding documentation for this env var on the webpage so that we're not causing surprise for users?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If we assume mkldnn is the only backend for kFComputeEx on cpu, shouldn't we make the change inside the caller of InferStorageType? Then we don't need to change it in every ops.
src/operator/nn/activation.cc
Outdated
@@ -116,6 +116,10 @@ inline static bool ActivationStorageType(const nnvm::NodeAttrs& attrs, | |||
if (dev_mask == mshadow::cpu::kDevMask && SupportMKLDNNAct(param)) { | |||
*dispatch_mode = DispatchMode::kFComputeEx; | |||
} | |||
if (!MKLDNNEnvSet()) { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Don't we need to check (dev_mask == mshadow::cpu::kDevMask) here?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looking good. Some notes:
- I suggest we align the env var name with the conventional prefix "MXNET_"
- As part of the PR, please also update env_var.md where all env vars are documented.
src/operator/nn/activation.cc
Outdated
@@ -158,6 +162,10 @@ inline static bool BackwardActStorageType(const nnvm::NodeAttrs& attrs, | |||
if (dev_mask == mshadow::cpu::kDevMask && SupportMKLDNNAct(param)) { | |||
*dispatch_mode = DispatchMode::kFComputeEx; | |||
} | |||
if (!MKLDNNEnvSet()) { | |||
*dispatch_mode = DispatchMode::kFComputeFallback; | |||
return ret; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Is this return statement redundant due to the statement in line 170?
@ZhennanQin there are some fcomputeex ops that are not mkldnn (customop is one example). something I did in another PR (#12019) was set an attr on all MKLDNN operators. if we wait for that merge we could then check if the operator is mkldnn and if the flag is set and fallback before calling inferstorage type. |
@pengzhao-intel review / approve please |
Looks good. One tip: Could we add an info message to let the user know the MKLDNN is disabled now by "MXNET_MKLDNN_ENABLED=0" and it can be switched on with env setting again? I afraid the user will forget this env setting in some bash file and then can't get the correct performance. |
No problem. Do you know the best place to put that info? I asked around last week and there doesn't seem to be a consensus. |
I look into the code. The message should be in the top level of OP but seem no proper place now. |
a3c94fd
to
3ad3ec3
Compare
Looks fine now even the fallback log is appeared many times :)
|
src/executor/attach_op_execs_pass.cc
Outdated
"You can re-enable by setting MXNET_MKLDNN_ENABLED=1"); | ||
#endif | ||
|
||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This warning only gets printed for symbolic. What about imperative/Gluon?
@@ -137,6 +137,10 @@ static inline bool SupportMKLDNN(const NDArray &input) { | |||
&& SupportStorageMKLDNN(input.storage_type()); | |||
} | |||
|
|||
static inline bool MKLDNNEnvSet() { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can we declare a static var so the linear time env lookup only happen once ?
3bc949f
to
e4bda35
Compare
e4bda35
to
65184d3
Compare
|
moved the log into the fallback code so it should work for gluon / imperative as well. that being said, we should probably clean up the logs in a future cause it's becoming messy. |
* fallback when env set * fix mkdnn default env * update docs to include desc about MXNET_MKLDNN_ENABLED * update env name to MXNET_MKLDNN_ENABLED * check dev_mask is cpu before fallback * log if flag is off * cache mkldnn check output * move logonce inside fallback * retrigger
Description
MKLDNN can be turned off with USE_MKLDNN env flag
Checklist
Essentials
Please feel free to remove inapplicable items for your PR.
Changes
Comments