add option to control use of MKL-DNN in jaxlib easyblock

[lexming]

A few enhancements to the jaxlib easyblock:

• use native CPU features for non-generic builds in EasyBuild
• add extra parameter to control the use of Intel's MKL-DNN
• only set support for NCCL starting from version 0.1.170

[boegel]

Test report by @boegel

Overview of tested easyconfigs (in order)

• SUCCESS jax-0.2.19-foss-2020b.eb
• SUCCESS jax-0.2.20-foss-2021a.eb

Build succeeded for 2 out of 2 (2 easyconfigs in total)
node3114.skitty.os - Linux centos linux 7.9.2009, x86_64, Intel(R) Xeon(R) Gold 6140 CPU @ 2.30GHz (skylake_avx512), Python 3.6.8
See https://gist.github.com/36ec3ca8e9beeb8a2a5332e370f3b83b for a full test report.

[branfosj]

Test report by @branfosj

Overview of tested easyconfigs (in order)

• SUCCESS jax-0.2.19-foss-2020b.eb
• SUCCESS jax-0.2.20-foss-2021a.eb

Build succeeded for 2 out of 2 (2 easyconfigs in total)
bear-pg0211u08b.bear.cluster - Linux RHEL 8.3, x86_64, Intel(R) Xeon(R) Gold 6248 CPU @ 2.50GHz (cascadelake), Python 3.6.8
See https://gist.github.com/74708d795783b7175a019711097ed305 for a full test report.

[branfosj]

lgtm

based on the tests here and in easybuilders/easybuild-easyconfigs#14334

[boegel]

lgtm

[boegel]

Test report by @boegel

Overview of tested easyconfigs (in order)

• SUCCESS jax-0.2.19-foss-2020b.eb
• SUCCESS jax-0.2.20-foss-2021a.eb

Build succeeded for 2 out of 2 (2 easyconfigs in total)
node3575.doduo.os - Linux RHEL 8.2, x86_64, AMD EPYC 7552 48-Core Processor (zen2), Python 3.6.8
See https://gist.github.com/43f07db6c9337848201353b6922633ba for a full test report.

[boegel]

Test report by @boegel

Overview of tested easyconfigs (in order)

• SUCCESS jax-0.2.19-foss-2020b.eb
• SUCCESS jax-0.2.20-foss-2021a.eb

Build succeeded for 2 out of 2 (2 easyconfigs in total)
node2624.swalot.os - Linux centos linux 7.9.2009, x86_64, Intel(R) Xeon(R) CPU E5-2660 v3 @ 2.60GHz (haswell), Python 3.6.8
See https://gist.github.com/a312005b01fce3d145010336a8a32d22 for a full test report.

[boegel]

Thanks @lexming!

[boegel]

Test report by @boegel

Overview of tested easyconfigs (in order)

• SUCCESS jax-0.2.19-fosscuda-2020b.eb

Build succeeded for 1 out of 1 (1 easyconfigs in total)
node3304.joltik.os - Linux centos linux 7.9.2009, x86_64, Intel(R) Xeon(R) Gold 6242 CPU @ 2.80GHz (cascadelake), 1 x NVIDIA NVIDIA Tesla V100-SXM2-32GB, 465.19.01, Python 3.6.8
See https://gist.github.com/96a5e7a85df29e81dcf6489b9c622472 for a full test report.