From d46db1bf6980c4d73d8af7a3a73c1f4b97113892 Mon Sep 17 00:00:00 2001
From: Andi Drebes <andi.drebes@zama.ai>
Date: Tue, 15 Nov 2022 16:49:50 +0100
Subject: [PATCH] feat(compiler): BConcrete: Add batched keyswitch and
 bootstrap

---
 .../Dialect/BConcrete/IR/BConcreteOps.td      | 26 +++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteOps.td b/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteOps.td
index 3763ed5f4..ebf0a7d05 100644
--- a/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteOps.td
+++ b/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteOps.td
@@ -85,6 +85,18 @@ def BConcrete_KeySwitchLweBufferOp : BConcrete_Op<"keyswitch_lwe_buffer"> {
     let results = (outs 1DTensorOf<[I64]>:$result);
 }
 
+def BConcrete_BatchedKeySwitchLweBufferOp : BConcrete_Op<"batched_keyswitch_lwe_buffer"> {
+    let arguments = (ins
+        // LweKeySwitchKeyType:$keyswitch_key,
+        2DTensorOf<[I64]>:$ciphertext,
+        I32Attr:$level,
+        I32Attr:$baseLog,
+        I32Attr:$lwe_dim_in,
+        I32Attr:$lwe_dim_out
+    );
+    let results = (outs 2DTensorOf<[I64]>:$result);
+}
+
 def BConcrete_BootstrapLweBufferOp : BConcrete_Op<"bootstrap_lwe_buffer"> {
     let arguments = (ins
         1DTensorOf<[I64]>:$input_ciphertext,
@@ -99,6 +111,20 @@ def BConcrete_BootstrapLweBufferOp : BConcrete_Op<"bootstrap_lwe_buffer"> {
     let results = (outs 1DTensorOf<[I64]>:$result);
 }
 
+def BConcrete_BatchedBootstrapLweBufferOp : BConcrete_Op<"batched_bootstrap_lwe_buffer"> {
+    let arguments = (ins
+        2DTensorOf<[I64]>:$input_ciphertext,
+        1DTensorOf<[I64]>:$lookup_table,
+        I32Attr:$inputLweDim,
+        I32Attr:$polySize,
+        I32Attr:$level,
+        I32Attr:$baseLog,
+        I32Attr:$glweDimension,
+        I32Attr:$outPrecision
+    );
+    let results = (outs 2DTensorOf<[I64]>:$result);
+}
+
 // TODO(16bits): hack
 def BConcrete_WopPBSCRTLweBufferOp : BConcrete_Op<"wop_pbs_crt_lwe_buffer"> {
     let arguments = (ins