diff --git a/.bazelrc b/.bazelrc index b5e1def120cc..fdcc9b4cec74 100644 --- a/.bazelrc +++ b/.bazelrc @@ -116,6 +116,7 @@ test --test_env=PJRT_LOCAL_PROCESS_RANK # This environmental variable is important for properly integrating with XLA. test --test_env=XLA_EXPERIMENTAL +test --test_env=XLA_STABLEHLO_COMPILE # To find `libpython` that is required to run tests (they run using installed wheels). test --test_env=LD_LIBRARY_PATH diff --git a/.github/scripts/run_tests.sh b/.github/scripts/run_tests.sh index 6d35cfb2dabe..a3704e618496 100755 --- a/.github/scripts/run_tests.sh +++ b/.github/scripts/run_tests.sh @@ -87,6 +87,8 @@ PYTORCH_DIR=$1 XLA_DIR=$2 USE_COVERAGE="${3:-0}" +export XLA_STABLEHLO_COMPILE=1 + if [ -x "$(command -v nvidia-smi)" ]; then num_devices=$(nvidia-smi --list-gpus | wc -l) echo "Found $num_devices GPU devices..." diff --git a/configuration.yaml b/configuration.yaml index c2f8c8290bd4..02f433d72e44 100644 --- a/configuration.yaml +++ b/configuration.yaml @@ -387,7 +387,7 @@ variables: flag is experimental. The default_value will be set to true when StableHLO workflow is mature. type: bool - default_value: false + default_value: true XLA_DUMP_POST_OPTIMIZATIONS: description: - Dump the HLO graph after optimizations. You need to use it together diff --git a/test/cpp/run_tests.sh b/test/cpp/run_tests.sh index d6b492dc6945..dffc8b709fc8 100755 --- a/test/cpp/run_tests.sh +++ b/test/cpp/run_tests.sh @@ -49,6 +49,7 @@ shift $(($OPTIND - 1)) # Set XLA_EXPERIMENTAL var to subsequently executed commands. export XLA_EXPERIMENTAL +export XLA_STABLEHLO_COMPILE=1 EXTRA_FLAGS="" diff --git a/test/neuron/run_tests.sh b/test/neuron/run_tests.sh index c5f277fe06a5..6ccf4422a01d 100755 --- a/test/neuron/run_tests.sh +++ b/test/neuron/run_tests.sh @@ -44,6 +44,7 @@ export TORCH_TEST_DEVICES="$CDIR/pytorch_test_base.py" export PYTORCH_TEST_WITH_SLOW=1 export XLA_DUMP_FATAL_STACK=1 export CPU_NUM_DEVICES=4 +export XLA_STABLEHLO_COMPILE=1 TORCH_XLA_DIR=$(cd ~; dirname "$(python -c 'import torch_xla; print(torch_xla.__file__)')") COVERAGE_FILE="$CDIR/../.coverage" diff --git a/test/run_tests.sh b/test/run_tests.sh index 142ca53f2ae2..f83896f54301 100755 --- a/test/run_tests.sh +++ b/test/run_tests.sh @@ -44,6 +44,7 @@ export TORCH_TEST_DEVICES="$CDIR/pytorch_test_base.py" export PYTORCH_TEST_WITH_SLOW=1 export XLA_DUMP_FATAL_STACK=1 export CPU_NUM_DEVICES=4 +unset XLA_STABLEHLO_COMPILE TORCH_XLA_DIR=$(cd ~; dirname "$(python -c 'import torch_xla; print(torch_xla.__file__)')") COVERAGE_FILE="$CDIR/../.coverage"