Skip to content
This repository has been archived by the owner on Nov 17, 2023. It is now read-only.

[BUGFIX] Pip Build: use runtime.Features instead of manual check for mkldnn headers (#19195) #19928

Merged
merged 1 commit into from
Feb 22, 2021
Merged
Changes from all commits
Commits
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
16 changes: 5 additions & 11 deletions tools/pip/setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,22 +142,16 @@ def skip_markdown_comments(md):
libraries.append('CUDA-10.0')
elif variant.startswith('CU92'):
libraries.append('CUDA-9.2')
elif variant.startswith('CU91'):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are we dropping support for all cuda till 9? I guess yes [since cuda11 has launched and we support till latest major version-1]. Can we update the PR title to also reflect the same? Since currently it just talks about mkldnn and this PR changes things more than that.

libraries.append('CUDA-9.1')
elif variant.startswith('CU90'):
libraries.append('CUDA-9.0')
elif variant.startswith('CU80'):
libraries.append('CUDA-8.0')
elif variant.startswith('CU75'):
libraries.append('CUDA-7.5')
if variant.endswith('MKL'):
libraries.append('MKLDNN')

from mxnet.runtime import Features
if Features().is_enabled("MKLDNN"):
libraries.append('MKLDNN')

short_description += ' This version uses {0}.'.format(' and '.join(libraries))

package_data = {'mxnet': [os.path.join('mxnet', os.path.basename(LIB_PATH[0]))],
'dmlc_tracker': []}
if variant != 'NATIVE':
if Features().is_enabled("MKLDNN"):
shutil.copytree(os.path.join(CURRENT_DIR, 'mxnet-build/3rdparty/mkldnn/include'),
os.path.join(CURRENT_DIR, 'mxnet/include/mkldnn'))
if platform.system() == 'Linux':
Expand Down