[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
diff --git a/.ci/pytorch/test.sh b/.ci/pytorch/test.sh
index 8fcb827..daf71c3 100755
--- a/.ci/pytorch/test.sh
+++ b/.ci/pytorch/test.sh
@@ -698,11 +698,14 @@
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
