[CI][dashboard] Add jemalloc back for aarch64 (#134189)

Forward fix based on https://github.com/pytorch/pytorch/pull/133997#discussion_r1726004220
Pull Request resolved: https://github.com/pytorch/pytorch/pull/134189
Approved by: https://github.com/malfet, https://github.com/huydhn
This commit is contained in:
Bin Bao 2024-08-22 21:08:39 +00:00 committed by PyTorch MergeBot
parent 1b6bbaa016
commit f0ba309d78

View File

@ -698,11 +698,14 @@ test_inductor_get_core_number() {
test_inductor_set_cpu_affinity(){
#set jemalloc
JEMALLOC_LIB="$(find /usr/lib -name libjemalloc.so.2)"
export LD_PRELOAD="$JEMALLOC_LIB":"$LD_PRELOAD"
export MALLOC_CONF="oversize_threshold:1,background_thread:true,metadata_thp:auto,dirty_decay_ms:-1,muzzy_decay_ms:-1"
if [[ "${TEST_CONFIG}" != *aarch64* ]]; then
JEMALLOC_LIB="$(find /usr/lib -name libjemalloc.so.2)"
# Use Intel OpenMP for x86
IOMP_LIB="$(dirname "$(which python)")/../lib/libiomp5.so"
export LD_PRELOAD="$JEMALLOC_LIB":"$IOMP_LIB":"$LD_PRELOAD"
export MALLOC_CONF="oversize_threshold:1,background_thread:true,metadata_thp:auto,dirty_decay_ms:-1,muzzy_decay_ms:-1"
export LD_PRELOAD="$IOMP_LIB":"$LD_PRELOAD"
export KMP_AFFINITY=granularity=fine,compact,1,0
export KMP_BLOCKTIME=1
fi