Skip to content

don't use --config=mkl for TensorFlow 2.4+ #2583

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Oct 27, 2021
Merged

don't use --config=mkl for TensorFlow 2.4+ #2583

Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
738b703
Don't use --config=mkl for TF 2.4+
Flamefire Sep 24, 2021
ba9b449
Add author
Flamefire Sep 27, 2021
572197b
clarify comment w.r.t. auto-enabling --config=mkl for TensorFlow vers…
boegel Oct 27, 2021
517129b
restructure handling of with_mkl_dnn logic in TensorFlow easyblock
boegel Oct 27, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 18 additions & 6 deletions easybuild/easyblocks/t/tensorflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
@author: Kenneth Hoste (HPC-UGent)
@author: Ake Sandgren (Umea University)
@author: Damian Alvarez (Forschungzentrum Juelich GmbH)
@author: Alexander Grund (TU Dresden)
"""
import glob
import os
Expand Down Expand Up @@ -213,13 +214,12 @@ class EB_TensorFlow(PythonPackage):

@staticmethod
def extra_options():
# We only want to install mkl-dnn by default on x86_64 systems
with_mkl_dnn_default = get_cpu_architecture() == X86_64
extra_vars = {
'path_filter': [[], "List of patterns to be filtered out in paths in $CPATH and $LIBRARY_PATH", CUSTOM],
'with_jemalloc': [None, "Make TensorFlow use jemalloc (usually enabled by default). " +
"Unsupported starting at TensorFlow 1.12!", CUSTOM],
'with_mkl_dnn': [with_mkl_dnn_default, "Make TensorFlow use Intel MKL-DNN", CUSTOM],
'with_mkl_dnn': [None, "Make TensorFlow use Intel MKL-DNN / oneDNN (enabled by default where supported)",
CUSTOM],
'with_xla': [None, "Enable XLA JIT compiler for possible runtime optimization of models", CUSTOM],
'test_script': [None, "Script to test TensorFlow installation with", CUSTOM],
'test_targets': [[], "List of Bazel targets which should be run during the test step", CUSTOM],
Expand Down Expand Up @@ -803,13 +803,25 @@ def build_step(self):
self.target_opts.append('--config=cuda')

# if mkl-dnn is listed as a dependency it is used. Otherwise downloaded if with_mkl_dnn is true
# Since 2.4.0 the mkl-config flag is not required as oneDNN is automatically used for x86 systems
# and MKL is no longer a dependency
mkl_root = get_software_root('mkl-dnn')
if mkl_root:
self.target_opts.append('--config=mkl')
env.setvar('TF_MKL_ROOT', mkl_root)
elif self.cfg['with_mkl_dnn']:
# this makes TensorFlow use mkl-dnn (cfr. https://github.com/01org/mkl-dnn)
self.target_opts.append('--config=mkl')
elif LooseVersion(self.version) < LooseVersion('2.4.0'):
# auto-enable use of MKL-DNN when possible if with_mkl_dnn is left unspecified
if self.cfg['with_mkl_dnn'] is None:
cpu_arch = get_cpu_architecture()
if cpu_arch == X86_64:
# Supported on x86 since forever
self.cfg['with_mkl_dnn'] = True
self.log.info("Auto-enabled use of MKL-DNN on %s CPU architecture", cpu_arch)
else:
self.log.info("Not enabling use of MKL-DNN on %s CPU architecture", cpu_arch)

if self.cfg['with_mkl_dnn']:
self.target_opts.append('--config=mkl')

# Compose final command
cmd = (
Expand Down