chore(hlapi): add cuda tests for FheBool

zama-ai · Mar 1, 2024 · 347fc9a · 347fc9a
1 parent 198485b
commit 347fc9a
Show file tree

Hide file tree

Showing 5 changed files with 335 additions and 73 deletions.
diff --git a/.github/workflows/aws_tfhe_gpu_tests.yml b/.github/workflows/aws_tfhe_gpu_tests.yml
@@ -112,6 +112,9 @@ jobs:
         run: |
           make test_c_api_gpu
 
+      - name: Run High Level API Tests
+        run: |
+          make test_high_level_api_gpu
 
       - name: Slack Notification
         if: ${{ always() }}

diff --git a/Makefile b/Makefile
@@ -533,6 +533,11 @@ test_high_level_api: install_rs_build_toolchain
 		--features=$(TARGET_ARCH_FEATURE),boolean,shortint,integer,internal-keycache -p $(TFHE_SPEC) \
 		-- high_level_api::
 
+test_high_level_api_gpu: install_rs_build_toolchain install_cargo_nextest
+	RUSTFLAGS="$(RUSTFLAGS)" cargo $(CARGO_RS_BUILD_TOOLCHAIN) nextest run --cargo-profile $(CARGO_PROFILE) \
+		--features=$(TARGET_ARCH_FEATURE),integer,internal-keycache,gpu -p $(TFHE_SPEC) \
+		-E "test(/high_level_api::.*gpu.*/)"
+
 .PHONY: test_user_doc # Run tests from the .md documentation
 test_user_doc: install_rs_build_toolchain
 	RUSTFLAGS="$(RUSTFLAGS)" cargo $(CARGO_RS_BUILD_TOOLCHAIN) test --profile $(CARGO_PROFILE) --doc \

diff --git a/tfhe/src/high_level_api/booleans/compact.rs b/tfhe/src/high_level_api/booleans/compact.rs
@@ -48,7 +48,9 @@ impl CompactFheBool {
         assert_eq!(ct.blocks.len(), 1);
         let mut block = BooleanBlock::new_unchecked(ct.blocks.into_iter().next().unwrap());
         block.0.degree = Degree::new(1);
-        FheBool::new(block)
+        let mut ciphertext = FheBool::new(block);
+        ciphertext.ciphertext.move_to_device_of_server_key_if_set();
+        ciphertext
     }
 }
 
@@ -131,7 +133,9 @@ impl CompactFheBoolList {
                 assert_eq!(ct.blocks.len(), 1);
                 let mut block = BooleanBlock::new_unchecked(ct.blocks.into_iter().next().unwrap());
                 block.0.degree = Degree::new(1);
-                FheBool::new(block)
+                let mut ciphertext = FheBool::new(block);
+                ciphertext.ciphertext.move_to_device_of_server_key_if_set();
+                ciphertext
             })
             .collect::<Vec<_>>()
     }

diff --git a/tfhe/src/high_level_api/booleans/compressed.rs b/tfhe/src/high_level_api/booleans/compressed.rs
@@ -43,9 +43,11 @@ impl CompressedFheBool {
     ///
     /// See [CompressedFheBool] example.
     pub fn decompress(&self) -> FheBool {
-        FheBool::new(BooleanBlock::new_unchecked(
+        let mut ciphertext = FheBool::new(BooleanBlock::new_unchecked(
             self.ciphertext.clone().decompress(),
-        ))
+        ));
+        ciphertext.ciphertext.move_to_device_of_server_key_if_set();
+        ciphertext
     }
 }