From 25362b2db20ca799436b14fae76aab8d5c5d4e15 Mon Sep 17 00:00:00 2001 From: Helder Campos Date: Fri, 30 May 2025 19:02:51 +0100 Subject: [PATCH] feat(hpu): Adding support for modulus switch mean compensation Including the pfail 2e-128 parameter set. Note: The HPU mockup still does not support mean compensation. --- .../config_store/sim/hpu_config.toml | 1 + .../sim/hpu_regif_core_cfg_1in3.toml | 15 +++++- .../config_store/u55c_gf64/hpu_config.toml | 1 + .../config_store/v80/hpu_config.toml | 1 + .../v80/hpu_regif_core_cfg_1in3.toml | 15 +++++- .../config_store/v80_archives/psi64.hpu | 4 +- .../tfhe-hpu-backend/src/interface/backend.rs | 19 ++++++++ .../tfhe-hpu-backend/src/interface/config.rs | 2 + .../src/interface/rtl/params.rs | 16 +++++++ .../src/interface/rtl/runtime.rs | 30 ++++++++++++ .../params/tuniform_64b_pfail128_psi64.toml | 48 +++++++++++++++++++ mockups/tfhe-hpu-mockup/src/lib.rs | 10 ++-- mockups/tfhe-hpu-mockup/src/modules/regmap.rs | 6 ++- 13 files changed, 157 insertions(+), 11 deletions(-) create mode 100644 mockups/tfhe-hpu-mockup/params/tuniform_64b_pfail128_psi64.toml diff --git a/backends/tfhe-hpu-backend/config_store/sim/hpu_config.toml b/backends/tfhe-hpu-backend/config_store/sim/hpu_config.toml index 80e7a4827..070dcf86b 100644 --- a/backends/tfhe-hpu-backend/config_store/sim/hpu_config.toml +++ b/backends/tfhe-hpu-backend/config_store/sim/hpu_config.toml @@ -13,6 +13,7 @@ bpip_use = true bpip_use_opportunism = true bpip_timeout = 100_000 + mod_switch_mean_comp = true [board] ct_mem = 32768 diff --git a/backends/tfhe-hpu-backend/config_store/sim/hpu_regif_core_cfg_1in3.toml b/backends/tfhe-hpu-backend/config_store/sim/hpu_regif_core_cfg_1in3.toml index bfdb80263..6b556f8fc 100644 --- a/backends/tfhe-hpu-backend/config_store/sim/hpu_regif_core_cfg_1in3.toml +++ b/backends/tfhe-hpu-backend/config_store/sim/hpu_regif_core_cfg_1in3.toml @@ -49,7 +49,8 @@ offset= 0x10 owner="Parameter" read_access="Read" write_access="None" - default={Param="VERSION"} + field.major={size_b=4, default={Param="VERSION_MAJOR"}, description="RTL major version"} + field.minor={size_b=4, default={Param="VERSION_MINOR"}, description="RTL minor version"} [section.info.register.ntt_architecture] description="NTT architecture" @@ -254,3 +255,15 @@ description="BPIP configuration" read_access="Read" write_access="Write" default={Cst=0xffffffff} + +# ===================================================================================================================== +[section.keyswitch] +offset= 0x3000 +description="Keyswitch Configuration" + +[section.keyswitch.register.config] + description="(1) Use use modulus switching mean compensation. (default), (0) Don't use modulus switching mean compensation." + owner="User" + read_access="Read" + write_access="Write" + field.mod_switch_mean_comp = { size_b=1, offset_b=0 , default={Cst=1}, description="Controls whether to use modulus switch mean compensation, aka. Mayeul's Trick."} diff --git a/backends/tfhe-hpu-backend/config_store/u55c_gf64/hpu_config.toml b/backends/tfhe-hpu-backend/config_store/u55c_gf64/hpu_config.toml index b5d460a1a..5d873d395 100644 --- a/backends/tfhe-hpu-backend/config_store/u55c_gf64/hpu_config.toml +++ b/backends/tfhe-hpu-backend/config_store/u55c_gf64/hpu_config.toml @@ -11,6 +11,7 @@ bpip_use = true bpip_use_opportunism = true bpip_timeout = 100_000 + mod_switch_mean_comp = true [board] ct_mem = 4096 diff --git a/backends/tfhe-hpu-backend/config_store/v80/hpu_config.toml b/backends/tfhe-hpu-backend/config_store/v80/hpu_config.toml index 8ed704322..6d00585f6 100644 --- a/backends/tfhe-hpu-backend/config_store/v80/hpu_config.toml +++ b/backends/tfhe-hpu-backend/config_store/v80/hpu_config.toml @@ -16,6 +16,7 @@ bpip_use = true bpip_use_opportunism = true bpip_timeout = 100_000 + mod_switch_mean_comp = true [board] ct_mem = 32768 diff --git a/backends/tfhe-hpu-backend/config_store/v80/hpu_regif_core_cfg_1in3.toml b/backends/tfhe-hpu-backend/config_store/v80/hpu_regif_core_cfg_1in3.toml index bfdb80263..6b556f8fc 100644 --- a/backends/tfhe-hpu-backend/config_store/v80/hpu_regif_core_cfg_1in3.toml +++ b/backends/tfhe-hpu-backend/config_store/v80/hpu_regif_core_cfg_1in3.toml @@ -49,7 +49,8 @@ offset= 0x10 owner="Parameter" read_access="Read" write_access="None" - default={Param="VERSION"} + field.major={size_b=4, default={Param="VERSION_MAJOR"}, description="RTL major version"} + field.minor={size_b=4, default={Param="VERSION_MINOR"}, description="RTL minor version"} [section.info.register.ntt_architecture] description="NTT architecture" @@ -254,3 +255,15 @@ description="BPIP configuration" read_access="Read" write_access="Write" default={Cst=0xffffffff} + +# ===================================================================================================================== +[section.keyswitch] +offset= 0x3000 +description="Keyswitch Configuration" + +[section.keyswitch.register.config] + description="(1) Use use modulus switching mean compensation. (default), (0) Don't use modulus switching mean compensation." + owner="User" + read_access="Read" + write_access="Write" + field.mod_switch_mean_comp = { size_b=1, offset_b=0 , default={Cst=1}, description="Controls whether to use modulus switch mean compensation, aka. Mayeul's Trick."} diff --git a/backends/tfhe-hpu-backend/config_store/v80_archives/psi64.hpu b/backends/tfhe-hpu-backend/config_store/v80_archives/psi64.hpu index c2add7619..cc82798da 100644 --- a/backends/tfhe-hpu-backend/config_store/v80_archives/psi64.hpu +++ b/backends/tfhe-hpu-backend/config_store/v80_archives/psi64.hpu @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:cb9ebedd0987130c4f6e1ef09f279d92f083815c1383da4b257198a33ab4881e -size 80293531 +oid sha256:4c17c71cedd183daeabf93649b61b294619445b695754606066051b10ecef27c +size 83001412 diff --git a/backends/tfhe-hpu-backend/src/interface/backend.rs b/backends/tfhe-hpu-backend/src/interface/backend.rs index ad38a6ec7..fa04697b8 100644 --- a/backends/tfhe-hpu-backend/src/interface/backend.rs +++ b/backends/tfhe-hpu-backend/src/interface/backend.rs @@ -165,7 +165,26 @@ impl HpuBackend { config.rtl.bpip_timeout, ); + let ks_config_reg = regmap + .register() + .get("keyswitch::config") + .expect("Unknown register, check regmap definition"); + hpu_hw.write_reg( + *ks_config_reg.offset() as u64, + ks_config_reg.from_field( + [( + "mod_switch_mean_comp", + config.rtl.mod_switch_mean_comp as u32, + )] + .into(), + ), + ); + info!("{params:?}"); + debug!( + "Keyswitch registers {:?}", + rtl::runtime::InfoKeyswitch::from_rtl(&mut hpu_hw, ®map) + ); debug!( "Isc registers {:?}", rtl::runtime::InfoIsc::from_rtl(&mut hpu_hw, ®map) diff --git a/backends/tfhe-hpu-backend/src/interface/config.rs b/backends/tfhe-hpu-backend/src/interface/config.rs index 822befb92..6e22f35a7 100644 --- a/backends/tfhe-hpu-backend/src/interface/config.rs +++ b/backends/tfhe-hpu-backend/src/interface/config.rs @@ -77,6 +77,8 @@ pub struct RtlConfig { pub bpip_use_opportunism: bool, /// Timeout value to start Bpip even if batch isn't full pub bpip_timeout: u32, + /// Use modulus switch mean compensation + pub mod_switch_mean_comp: bool, } /// On-board memory configuration diff --git a/backends/tfhe-hpu-backend/src/interface/rtl/params.rs b/backends/tfhe-hpu-backend/src/interface/rtl/params.rs index 4cf45db69..a2c68253e 100644 --- a/backends/tfhe-hpu-backend/src/interface/rtl/params.rs +++ b/backends/tfhe-hpu-backend/src/interface/rtl/params.rs @@ -416,6 +416,21 @@ pub const MSG2_CARRY2_PFAIL64_132B_TUNIFORM_7E47D8C: HpuPBSParameters = HpuPBSPa ciphertext_width: 64, }; +pub const MSG2_CARRY2_PFAIL128_132B_TUNIFORM_144A47: HpuPBSParameters = HpuPBSParameters { + lwe_dimension: 879, + glwe_dimension: 1, + polynomial_size: 2048, + lwe_noise_distribution: HpuNoiseDistributionInput::TUniformBound(3), + glwe_noise_distribution: HpuNoiseDistributionInput::TUniformBound(17), + pbs_base_log: 23, + pbs_level: 1, + ks_base_log: 2, + ks_level: 8, + message_width: 2, + carry_width: 2, + ciphertext_width: 64, +}; + impl FromRtl for HpuPBSParameters { fn from_rtl(ffi_hw: &mut ffi::HpuHw, regmap: &FlatRegmap) -> Self { let pbs_app = regmap @@ -456,6 +471,7 @@ impl FromRtl for HpuPBSParameters { 11 => MSG2_CARRY2_TUNIFORM, 12 => MSG2_CARRY2_PFAIL64_132B_GAUSSIAN_1F72DBA, 13 => MSG2_CARRY2_PFAIL64_132B_TUNIFORM_7E47D8C, + 14 => MSG2_CARRY2_PFAIL128_132B_TUNIFORM_144A47, _ => panic!("Unknown TfheAppName encoding"), } } diff --git a/backends/tfhe-hpu-backend/src/interface/rtl/runtime.rs b/backends/tfhe-hpu-backend/src/interface/rtl/runtime.rs index 3e54ebaeb..cfe74210e 100644 --- a/backends/tfhe-hpu-backend/src/interface/rtl/runtime.rs +++ b/backends/tfhe-hpu-backend/src/interface/rtl/runtime.rs @@ -932,3 +932,33 @@ impl ErrorHpu { self.error_3in3 = ffi_hw.read_reg(*reg.offset() as u64); } } + +#[derive(Debug, Default)] +pub struct InfoKeyswitch { + /// Use modulus switch mean compensation + mod_switch_mean_comp: bool, +} + +impl FromRtl for InfoKeyswitch { + fn from_rtl(ffi_hw: &mut ffi::HpuHw, regmap: &FlatRegmap) -> Self { + // Info structure have method to update + // Instead of redefine parsing here, use a default construct and update methods + let mut infos = Self::default(); + infos.update(ffi_hw, regmap); + infos + } +} + +impl InfoKeyswitch { + pub fn update_mod_switch_mean_comp(&mut self, ffi_hw: &mut ffi::HpuHw, regmap: &FlatRegmap) { + let reg = regmap + .register() + .get("keyswitch::config") + .expect("Unknown register, check regmap definition"); + self.mod_switch_mean_comp = ffi_hw.read_reg(*reg.offset() as u64) != 0; + } + + pub fn update(&mut self, ffi_hw: &mut ffi::HpuHw, regmap: &FlatRegmap) { + self.update_mod_switch_mean_comp(ffi_hw, regmap); + } +} diff --git a/mockups/tfhe-hpu-mockup/params/tuniform_64b_pfail128_psi64.toml b/mockups/tfhe-hpu-mockup/params/tuniform_64b_pfail128_psi64.toml new file mode 100644 index 000000000..aa05e15a9 --- /dev/null +++ b/mockups/tfhe-hpu-mockup/params/tuniform_64b_pfail128_psi64.toml @@ -0,0 +1,48 @@ +[pbs_params] + lwe_dimension=879 + glwe_dimension=1 + polynomial_size=2048 + lwe_noise_distribution={TUniformBound= 3} + glwe_noise_distribution={TUniformBound= 17} + pbs_base_log= 23 + pbs_level= 1 + ks_base_log= 2 + ks_level= 8 + message_width= 2 + carry_width= 2 + ciphertext_width= 64 + opportunistic=true + +[ntt_params] + core_arch= {GF64=[5,6]} + min_pbs_nb= 11 + batch_pbs_nb= 12 + total_pbs_nb= 32 + ct_width= 64 + radix= 2 + stg_nb= 11 + prime_modulus= "GF64" + psi= 64 + delta= 5 + +[ks_params] + width= 21 + lbx= 3 + lby= 64 + lbz= 3 + +[pc_params] + ksk_pc= 16 + ksk_bytes_w= 32 + bsk_pc= 8 + bsk_bytes_w= 32 + pem_pc= 2 + pem_bytes_w= 32 + glwe_bytes_w= 32 + +[regf_params] + reg_nb= 64 + coef_nb= 32 +[isc_params] + min_iop_size= 4 + depth= 64 diff --git a/mockups/tfhe-hpu-mockup/src/lib.rs b/mockups/tfhe-hpu-mockup/src/lib.rs index e37bf6cdb..c01540297 100644 --- a/mockups/tfhe-hpu-mockup/src/lib.rs +++ b/mockups/tfhe-hpu-mockup/src/lib.rs @@ -29,8 +29,6 @@ use modules::{DdrMem, HbmBank, RegisterEvent, RegisterMap, UCore, HBM_BANK_NB}; use tfhe::tfhe_hpu_backend::interface::io_dump::HexMem; use tfhe::tfhe_hpu_backend::prelude::*; -use serde_json; - pub struct HpuSim { config: HpuConfig, params: MockupParameters, @@ -178,7 +176,7 @@ impl HpuSim { } RegisterReq::PbsParams => { self.ipc.register_ack(RegisterAck::PbsParams( - self.params.rtl_params.pbs_params.clone(), + self.params.rtl_params.pbs_params, )); } } @@ -672,7 +670,7 @@ impl HpuSim { // Compute Lut properties let (modulus_sup, box_size, fn_stride) = { let pbs_p = &self.params.rtl_params.pbs_params; - let modulus_sup = 1_usize << pbs_p.message_width + pbs_p.carry_width; + let modulus_sup = 1_usize << (pbs_p.message_width + pbs_p.carry_width); let box_size = pbs_p.polynomial_size / modulus_sup; // Max valid degree for a ciphertext when using the LUT we generate // If MaxDegree == 1, we can have two input values 0 and 1, so we need MaxDegree + 1 @@ -725,7 +723,7 @@ impl HpuSim { let bfr_after_ms = lwe_ciphertext_modulus_switch(bfr_after_ks.as_view(), log_modulus); - blind_rotate_ntt64_bnf_assign(&bfr_after_ms, &mut tfhe_lut, &bsk); + blind_rotate_ntt64_bnf_assign(&bfr_after_ms, &mut tfhe_lut, bsk); assert_eq!( dst_rid.0, @@ -736,7 +734,7 @@ impl HpuSim { // Compute ManyLut function stride let fn_stride = { let pbs_p = &self.params.rtl_params.pbs_params; - let modulus_sup = 1_usize << pbs_p.message_width + pbs_p.carry_width; + let modulus_sup = 1_usize << (pbs_p.message_width + pbs_p.carry_width); let box_size = pbs_p.polynomial_size / modulus_sup; // Max valid degree for a ciphertext when using the LUT we generate // If MaxDegree == 1, we can have two input values 0 and 1, so we need MaxDegree + 1 diff --git a/mockups/tfhe-hpu-mockup/src/modules/regmap.rs b/mockups/tfhe-hpu-mockup/src/modules/regmap.rs index c08bd8487..f6ff87016 100644 --- a/mockups/tfhe-hpu-mockup/src/modules/regmap.rs +++ b/mockups/tfhe-hpu-mockup/src/modules/regmap.rs @@ -127,7 +127,7 @@ impl RegisterMap { (ntt_p.batch_pbs_nb + (ntt_p.total_pbs_nb << 8)) as u32 } "info::ntt_modulo" => { - MOD_NTT_NAME_OFS + (self.rtl_params.ntt_params.prime_modulus.clone() as u8) as u32 + MOD_NTT_NAME_OFS + (self.rtl_params.ntt_params.prime_modulus as u8) as u32 } "info::application" => { @@ -147,6 +147,10 @@ impl RegisterMap { APPLICATION_NAME_OFS + 11 } else if MSG2_CARRY2_PFAIL64_132B_GAUSSIAN_1F72DBA == self.rtl_params.pbs_params { APPLICATION_NAME_OFS + 12 + } else if MSG2_CARRY2_PFAIL64_132B_TUNIFORM_7E47D8C == self.rtl_params.pbs_params { + APPLICATION_NAME_OFS + 13 + } else if MSG2_CARRY2_PFAIL128_132B_TUNIFORM_144A47 == self.rtl_params.pbs_params { + APPLICATION_NAME_OFS + 14 } else { // Custom simulation parameters set // -> Return 1 without NAME_OFS