chore(gpu): add better panic when calling old compression functions

agnesLeroy · agnesLeroy · commit 5968778b5752 · 2025-11-20T14:29:11.000+01:00
diff --git a/tfhe/src/high_level_api/booleans/compressed.rs b/tfhe/src/high_level_api/booleans/compressed.rs
@@ -2,7 +2,8 @@ use crate::backward_compatibility::booleans::{
     CompressedFheBoolVersions, InnerCompressedFheBoolVersions,
 };
 use crate::conformance::ParameterSetConformant;
-use crate::high_level_api::global_state::with_cpu_internal_keys;
+use crate::high_level_api::global_state;
+use crate::high_level_api::keys::InternalServerKey;
 use crate::high_level_api::re_randomization::ReRandomizationMetadata;
 use crate::high_level_api::traits::Tagged;
 use crate::integer::BooleanBlock;
@@ -74,17 +75,27 @@ impl CompressedFheBool {
         let ciphertext = BooleanBlock::new_unchecked(match &self.inner {
             InnerCompressedFheBool::Seeded(seeded) => seeded.decompress(),
             InnerCompressedFheBool::ModulusSwitched(modulus_switched) => {
-                with_cpu_internal_keys(|sk| sk.pbs_key().key.decompress(modulus_switched))
+                global_state::with_internal_keys(|keys| match keys {
+                    InternalServerKey::Cpu(cpu_key) => {
+                        cpu_key.pbs_key().key.decompress(modulus_switched)
+                    }
+                    #[cfg(feature = "gpu")]
+                    InternalServerKey::Cuda(_) => {
+                        panic!("decompress() on FheBool is not supported on GPU, use a CompressedCiphertextList instead");
+                    }
+                    #[cfg(feature = "hpu")]
+                    InternalServerKey::Hpu(_) => {
+                        panic!("decompress() on FheBool is not supported on HPU devices");
+                    }
+                })
             }
         });
         let mut ciphertext = FheBool::new(
             ciphertext,
             self.tag.clone(),
             ReRandomizationMetadata::default(),
         );
-
         ciphertext.ciphertext.move_to_device_of_server_key_if_set();
-
         ciphertext
     }
 }
@@ -117,16 +128,26 @@ impl Named for CompressedFheBool {
 
 impl FheBool {
     pub fn compress(&self) -> CompressedFheBool {
-        with_cpu_internal_keys(|sk| {
-            let inner = InnerCompressedFheBool::ModulusSwitched(
-                sk.pbs_key()
-                    .key
-                    .switch_modulus_and_compress(&self.ciphertext.on_cpu().0),
-            );
-
-            CompressedFheBool {
-                inner,
-                tag: sk.tag.clone(),
+        global_state::with_internal_keys(|keys| match keys {
+            InternalServerKey::Cpu(cpu_key) => {
+                let inner = InnerCompressedFheBool::ModulusSwitched(
+                    cpu_key
+                        .pbs_key()
+                        .key
+                        .switch_modulus_and_compress(&self.ciphertext.on_cpu().0),
+                );
+                CompressedFheBool {
+                    inner,
+                    tag: cpu_key.tag.clone(),
+                }
+            }
+            #[cfg(feature = "gpu")]
+            InternalServerKey::Cuda(_) => {
+                panic!("compress() on FheBool is not supported on GPU, use a CompressedCiphertextList instead");
+            }
+            #[cfg(feature = "hpu")]
+            InternalServerKey::Hpu(_) => {
+                panic!("compress() on FheBool is not supported on HPU devices");
             }
         })
     }
diff --git a/tfhe/src/high_level_api/integers/signed/compressed.rs b/tfhe/src/high_level_api/integers/signed/compressed.rs
@@ -5,9 +5,10 @@ use crate::backward_compatibility::integers::{
 };
 use crate::conformance::ParameterSetConformant;
 use crate::core_crypto::prelude::SignedNumeric;
-use crate::high_level_api::global_state::with_cpu_internal_keys;
+use crate::high_level_api::global_state;
 use crate::high_level_api::integers::signed::base::FheIntConformanceParams;
 use crate::high_level_api::integers::{FheInt, FheIntId};
+use crate::high_level_api::keys::InternalServerKey;
 use crate::high_level_api::re_randomization::ReRandomizationMetadata;
 use crate::high_level_api::traits::Tagged;
 use crate::integer::block_decomposition::DecomposableInto;
@@ -111,7 +112,19 @@ where
         let ciphertext = match &self.ciphertext {
             CompressedSignedRadixCiphertext::Seeded(ct) => ct.decompress(),
             CompressedSignedRadixCiphertext::ModulusSwitched(ct) => {
-                with_cpu_internal_keys(|sk| sk.pbs_key().decompress_signed_parallelized(ct))
+                global_state::with_internal_keys(|keys| match keys {
+                    InternalServerKey::Cpu(cpu_key) => {
+                        cpu_key.pbs_key().decompress_signed_parallelized(ct)
+                    }
+                    #[cfg(feature = "gpu")]
+                    InternalServerKey::Cuda(_) => {
+                        panic!("decompress() on FheInt is not supported on GPU, use a CompressedCiphertextList instead");
+                    }
+                    #[cfg(feature = "hpu")]
+                    InternalServerKey::Hpu(_) => {
+                        panic!("decompress() on FheInt is not supported on HPU devices");
+                    }
+                })
             }
         };
         FheInt::new(
@@ -180,14 +193,25 @@ where
     Id: FheIntId,
 {
     pub fn compress(&self) -> CompressedFheInt<Id> {
-        let a = with_cpu_internal_keys(|sk| {
-            sk.pbs_key()
-                .switch_modulus_and_compress_signed_parallelized(&self.ciphertext.on_cpu())
-        });
-
-        CompressedFheInt::new(
-            CompressedSignedRadixCiphertext::ModulusSwitched(a),
-            self.tag.clone(),
-        )
+        global_state::with_internal_keys(|keys| match keys {
+            InternalServerKey::Cpu(cpu_key) => {
+                let a = cpu_key
+                    .pbs_key()
+                    .switch_modulus_and_compress_signed_parallelized(&self.ciphertext.on_cpu());
+
+                CompressedFheInt::new(
+                    CompressedSignedRadixCiphertext::ModulusSwitched(a),
+                    self.tag.clone(),
+                )
+            }
+            #[cfg(feature = "gpu")]
+            InternalServerKey::Cuda(_) => {
+                panic!("compress() on FheInt is not supported on GPU, use a CompressedCiphertextList instead");
+            }
+            #[cfg(feature = "hpu")]
+            InternalServerKey::Hpu(_) => {
+                panic!("compress() on FheInt is not supported on HPU devices");
+            }
+        })
     }
 }
diff --git a/tfhe/src/high_level_api/integers/unsigned/compressed.rs b/tfhe/src/high_level_api/integers/unsigned/compressed.rs
@@ -5,13 +5,13 @@ use crate::backward_compatibility::integers::{
 };
 use crate::conformance::ParameterSetConformant;
 use crate::core_crypto::prelude::UnsignedNumeric;
-use crate::high_level_api::global_state::with_cpu_internal_keys;
 use crate::high_level_api::integers::unsigned::base::{
     FheUint, FheUintConformanceParams, FheUintId,
 };
+use crate::high_level_api::keys::InternalServerKey;
 use crate::high_level_api::re_randomization::ReRandomizationMetadata;
 use crate::high_level_api::traits::{FheTryEncrypt, Tagged};
-use crate::high_level_api::ClientKey;
+use crate::high_level_api::{global_state, ClientKey};
 use crate::integer::block_decomposition::DecomposableInto;
 use crate::integer::ciphertext::{
     CompressedModulusSwitchedRadixCiphertext,
@@ -108,7 +108,19 @@ where
         let inner = match &self.ciphertext {
             CompressedRadixCiphertext::Seeded(ct) => ct.decompress(),
             CompressedRadixCiphertext::ModulusSwitched(ct) => {
-                with_cpu_internal_keys(|sk| sk.pbs_key().decompress_parallelized(ct))
+                global_state::with_internal_keys(|keys| match keys {
+                    InternalServerKey::Cpu(cpu_key) => {
+                        cpu_key.pbs_key().decompress_parallelized(ct)
+                    }
+                    #[cfg(feature = "gpu")]
+                    InternalServerKey::Cuda(_) => {
+                        panic!("decompress() on FheUint is not supported on GPU, use a CompressedCiphertextList instead");
+                    }
+                    #[cfg(feature = "hpu")]
+                    InternalServerKey::Hpu(_) => {
+                        panic!("decompress() on FheUint is not supported on HPU devices");
+                    }
+                })
             }
         };
 
@@ -179,11 +191,24 @@ where
     Id: FheUintId,
 {
     pub fn compress(&self) -> CompressedFheUint<Id> {
-        let ciphertext = CompressedRadixCiphertext::ModulusSwitched(with_cpu_internal_keys(|sk| {
-            sk.pbs_key()
-                .switch_modulus_and_compress_parallelized(&self.ciphertext.on_cpu())
-        }));
-        CompressedFheUint::new(ciphertext, self.tag.clone())
+        global_state::with_internal_keys(|keys| match keys {
+            InternalServerKey::Cpu(cpu_key) => {
+                let ciphertext = CompressedRadixCiphertext::ModulusSwitched(
+                    cpu_key
+                        .pbs_key()
+                        .switch_modulus_and_compress_parallelized(&self.ciphertext.on_cpu()),
+                );
+                CompressedFheUint::new(ciphertext, self.tag.clone())
+            }
+            #[cfg(feature = "gpu")]
+            InternalServerKey::Cuda(_) => {
+                panic!("compress() on FheUint is not supported on GPU, use a CompressedCiphertextList instead");
+            }
+            #[cfg(feature = "hpu")]
+            InternalServerKey::Hpu(_) => {
+                panic!("compress() on FheUint is not supported on HPU devices");
+            }
+        })
     }
 }
 
