chore(ci): enable AVX512 for integer and multi bit integer tests

zama-ai · Oct 9, 2023 · 140ac2e · 140ac2e
1 parent bb3c8e7
commit 140ac2e
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 10 deletions.
diff --git a/.github/workflows/aws_tfhe_integer_tests.yml b/.github/workflows/aws_tfhe_integer_tests.yml
@@ -4,6 +4,7 @@ env:
   CARGO_TERM_COLOR: always
   ACTION_RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
   RUSTFLAGS: "-C target-cpu=native"
+  AVX512_SUPPORT: "ON"
 
 on:
   # Allows you to run this workflow manually from the Actions tab as an alternative.

diff --git a/.github/workflows/aws_tfhe_multi_bit_tests.yml b/.github/workflows/aws_tfhe_multi_bit_tests.yml
@@ -4,6 +4,7 @@ env:
   CARGO_TERM_COLOR: always
   ACTION_RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
   RUSTFLAGS: "-C target-cpu=native"
+  AVX512_SUPPORT: "ON"
 
 on:
   # Allows you to run this workflow manually from the Actions tab as an alternative.

diff --git a/Makefile b/Makefile
@@ -363,18 +363,18 @@ test_shortint_cov: install_rs_check_toolchain install_tarpaulin
 		-p tfhe -- shortint::
 
 .PHONY: test_integer_ci # Run the tests for integer ci
-test_integer_ci: install_rs_build_toolchain install_cargo_nextest
+test_integer_ci: install_rs_check_toolchain install_cargo_nextest
 	BIG_TESTS_INSTANCE="$(BIG_TESTS_INSTANCE)" \
 	FAST_TESTS="$(FAST_TESTS)" \
-		./scripts/integer-tests.sh --rust-toolchain $(CARGO_RS_BUILD_TOOLCHAIN) \
-		--cargo-profile "$(CARGO_PROFILE)"
+		./scripts/integer-tests.sh --rust-toolchain $(CARGO_RS_CHECK_TOOLCHAIN) \
+		--cargo-profile "$(CARGO_PROFILE)" --avx512-support "$(AVX512_SUPPORT)"
 
 .PHONY: test_integer_multi_bit_ci # Run the tests for integer ci running only multibit tests
-test_integer_multi_bit_ci: install_rs_build_toolchain install_cargo_nextest
+test_integer_multi_bit_ci: install_rs_check_toolchain install_cargo_nextest
 	BIG_TESTS_INSTANCE="$(BIG_TESTS_INSTANCE)" \
 	FAST_TESTS="$(FAST_TESTS)" \
-		./scripts/integer-tests.sh --rust-toolchain $(CARGO_RS_BUILD_TOOLCHAIN) \
-		--cargo-profile "$(CARGO_PROFILE)" --multi-bit
+		./scripts/integer-tests.sh --rust-toolchain $(CARGO_RS_CHECK_TOOLCHAIN) \
+		--cargo-profile "$(CARGO_PROFILE)" --multi-bit --avx512-support "$(AVX512_SUPPORT)"
 
 .PHONY: test_safe_deserialization # Run the tests for safe deserialization
 test_safe_deserialization: install_rs_build_toolchain install_cargo_nextest

diff --git a/scripts/integer-tests.sh b/scripts/integer-tests.sh
@@ -9,13 +9,15 @@ function usage() {
     echo "--rust-toolchain          The toolchain to run the tests with default: stable"
     echo "--multi-bit               Run multi-bit tests only: default off"
     echo "--cargo-profile           The cargo profile used to build tests"
+    echo "--avx512-support          Set to ON to enable avx512"
     echo
 }
 
 RUST_TOOLCHAIN="+stable"
 multi_bit=""
 not_multi_bit="_multi_bit"
 cargo_profile="release"
+avx512_feature=""
 
 while [ -n "$1" ]
 do
@@ -40,6 +42,13 @@ do
             cargo_profile="$1"
             ;;
 
+        "--avx512-support" )
+            shift
+            if [[ "$1" == "ON" ]]; then
+                avx512_feature=nightly-avx512
+            fi
+            ;;
+
         *)
             echo "Unknown param : $1"
             exit 1
@@ -104,15 +113,15 @@ and not test(/.*default_add_sequence_multi_thread_param_message_3_carry_3_ks_pbs
         --cargo-profile "${cargo_profile}" \
         --package tfhe \
         --profile ci \
-        --features="${ARCH_FEATURE}",integer,internal-keycache \
+        --features="${ARCH_FEATURE}",integer,internal-keycache,"${avx512_feature}" \
         --test-threads "${n_threads}" \
         -E "$filter_expression"
 
     if [[ "${multi_bit}" == "" ]]; then
         cargo "${RUST_TOOLCHAIN}" test \
             --profile "${cargo_profile}" \
             --package tfhe \
-            --features="${ARCH_FEATURE}",integer,internal-keycache \
+            --features="${ARCH_FEATURE}",integer,internal-keycache,"${avx512_feature}" \
             --doc \
             -- integer::
     fi
@@ -148,15 +157,15 @@ and not test(/.*default_add_sequence_multi_thread_param_message_3_carry_3_ks_pbs
         --cargo-profile "${cargo_profile}" \
         --package tfhe \
         --profile ci \
-        --features="${ARCH_FEATURE}",integer,internal-keycache \
+        --features="${ARCH_FEATURE}",integer,internal-keycache,"${avx512_feature}" \
         --test-threads $num_threads \
         -E "$filter_expression"
 
     if [[ "${multi_bit}" == "" ]]; then
         cargo "${RUST_TOOLCHAIN}" test \
             --profile "${cargo_profile}" \
             --package tfhe \
-            --features="${ARCH_FEATURE}",integer,internal-keycache \
+            --features="${ARCH_FEATURE}",integer,internal-keycache,"${avx512_feature}" \
             --doc \
             -- --test-threads="$(${nproc_bin})" integer::
     fi