fix(compiler): add missing args in woppbs simulation

zama-ai · Oct 24, 2023 · 9b88dab · 9b88dab
1 parent 4cebe95
commit 9b88dab
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 5 deletions.
diff --git a/compilers/concrete-compiler/compiler/include/concretelang/Runtime/simulation.h b/compilers/concrete-compiler/compiler/include/concretelang/Runtime/simulation.h
@@ -76,7 +76,8 @@ void sim_wop_pbs_crt(
     // Additional crypto parameters
     uint32_t lwe_small_dim, uint32_t cbs_level_count, uint32_t cbs_base_log,
     uint32_t ksk_level_count, uint32_t ksk_base_log, uint32_t bsk_level_count,
-    uint32_t bsk_base_log, uint32_t polynomial_size, uint32_t glwe_dim);
+    uint32_t bsk_base_log, uint32_t polynomial_size, uint32_t pksk_base_log,
+    uint32_t pksk_level_count, uint32_t glwe_dim);
 
 void sim_encode_expand_lut_for_boostrap(
     uint64_t *in_allocated, uint64_t *in_aligned, uint64_t in_offset,

diff --git a/compilers/concrete-compiler/compiler/lib/Conversion/SimulateTFHE/SimulateTFHE.cpp b/compilers/concrete-compiler/compiler/lib/Conversion/SimulateTFHE/SimulateTFHE.cpp
@@ -354,6 +354,10 @@ struct WopPBSGLWEOpPattern
         wopPbs.getLoc(), adaptor.getBsk().getLevels(), 32);
     auto bskBaseLogCst = rewriter.create<mlir::arith::ConstantIntOp>(
         wopPbs.getLoc(), adaptor.getBsk().getBaseLog(), 32);
+    auto fpkskLevelCountCst = rewriter.create<mlir::arith::ConstantIntOp>(
+        wopPbs.getLoc(), adaptor.getPksk().getLevels(), 32);
+    auto fpkskBaseLogCst = rewriter.create<mlir::arith::ConstantIntOp>(
+        wopPbs.getLoc(), adaptor.getPksk().getBaseLog(), 32);
     auto polySizeCst = rewriter.create<mlir::arith::ConstantIntOp>(
         wopPbs.getLoc(), adaptor.getPksk().getOutputPolySize(), 32);
     auto glweDimCst = rewriter.create<mlir::arith::ConstantIntOp>(
@@ -370,6 +374,7 @@ struct WopPBSGLWEOpPattern
                  rewriter.getIntegerType(32), rewriter.getIntegerType(32),
                  rewriter.getIntegerType(32), rewriter.getIntegerType(32),
                  rewriter.getIntegerType(32), rewriter.getIntegerType(32),
+                 rewriter.getIntegerType(32), rewriter.getIntegerType(32),
                  rewriter.getIntegerType(32), rewriter.getIntegerType(32)},
                 {}))
             .failed()) {
@@ -382,7 +387,7 @@ struct WopPBSGLWEOpPattern
                           crtDecompValue, lweDimCst, cbsLevelCountCst,
                           cbsBaseLogCst, kskLevelCountCst, kskBaseLogCst,
                           bskLevelCountCst, bskBaseLogCst, polySizeCst,
-                          glweDimCst}));
+                          fpkskLevelCountCst, fpkskBaseLogCst, glweDimCst}));
 
     rewriter.replaceOp(wopPbs, outputBuffer);
 

diff --git a/compilers/concrete-compiler/compiler/lib/Runtime/simulation.cpp b/compilers/concrete-compiler/compiler/lib/Runtime/simulation.cpp
@@ -111,7 +111,8 @@ void sim_wop_pbs_crt(
     // Additional crypto parameters
     uint32_t lwe_small_dim, uint32_t cbs_level_count, uint32_t cbs_base_log,
     uint32_t ksk_level_count, uint32_t ksk_base_log, uint32_t bsk_level_count,
-    uint32_t bsk_base_log, uint32_t polynomial_size, uint32_t glwe_dim) {
+    uint32_t bsk_base_log, uint32_t polynomial_size, uint32_t pksk_base_log,
+    uint32_t pksk_level_count, uint32_t glwe_dim) {
 
   // Check number of blocks
   assert(out_size == in_size && out_size == crt_decomp_size);
@@ -172,8 +173,8 @@ void sim_wop_pbs_crt(
   simulation_circuit_bootstrap_boolean_vertical_packing_lwe_ciphertext_u64(
       extract_bits_output_buffer, out_aligned + out_offset, ct_in_count,
       ct_out_count, lut_size, lut_count, lut_ct_aligned + lut_ct_offset,
-      glwe_dim, log_poly_size, lwe_small_dim, cbs_level_count, cbs_base_log, 64,
-      128);
+      glwe_dim, log_poly_size, lwe_small_dim, bsk_level_count, bsk_base_log,
+      cbs_level_count, cbs_base_log, pksk_level_count, pksk_base_log, 64, 128);
 }
 
 uint64_t sim_neg_lwe_u64(uint64_t plaintext) { return ~plaintext + 1; }