chore(hlapi): add tests booleans on GPU

zama-ai · Feb 29, 2024 · fab719c · fab719c
1 parent 46a87c6
commit fab719c
Show file tree

Hide file tree

Showing 6 changed files with 349 additions and 91 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
+	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/base.rs b/tfhe/src/high_level_api/booleans/base.rs
@@ -325,15 +325,13 @@ impl FheEq<bool> for FheBool {
                 InnerBoolean::Cpu(BooleanBlock::new_unchecked(inner))
             }
             #[cfg(feature = "gpu")]
-            InternalServerKey::Cuda(_) => {
-                // with_thread_local_cuda_stream(|stream| {
-                //                let inner =
-                //                    cuda_key
-                //                        .key
-                //                        .scalar_eq(&self.ciphertext.on_gpu(), u8::from(other),
-                // stream);                InnerBoolean::Cuda(inner)
-                todo!()
-            }
+            InternalServerKey::Cuda(cuda_key) => with_thread_local_cuda_stream(|stream| {
+                let inner =
+                    cuda_key
+                        .key
+                        .scalar_eq(&self.ciphertext.on_gpu(), u8::from(other), stream);
+                InnerBoolean::Cuda(inner)
+            }),
         });
         Self::new(ciphertext)
     }
@@ -366,15 +364,13 @@ impl FheEq<bool> for FheBool {
                 InnerBoolean::Cpu(BooleanBlock::new_unchecked(inner))
             }
             #[cfg(feature = "gpu")]
-            InternalServerKey::Cuda(_) => {
-                //            with_thread_local_cuda_stream(|stream| {
-                //                let inner =
-                //                    cuda_key
-                //                        .key
-                //                        .scalar_ne(&self.ciphertext.on_gpu(), u8::from(other),
-                // stream);                InnerBoolean::Cuda(inner)
-                todo!()
-            }
+            InternalServerKey::Cuda(cuda_key) => with_thread_local_cuda_stream(|stream| {
+                let inner =
+                    cuda_key
+                        .key
+                        .scalar_ne(&self.ciphertext.on_gpu(), u8::from(other), stream);
+                InnerBoolean::Cuda(inner)
+            }),
         });
         Self::new(ciphertext)
     }

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
     }
 }