fix(compiler): Fix tfhe global parametrization to handle k>1

Co-authored-by: Mayeul@Zama <mayeul.debellabre@zama.ai>
2026-02-09 03:55:04 -05:00 · 2022-06-23 15:21:10 +02:00
parent ef9d11c16f
commit 2de76e9c4e
22 changed files with 228 additions and 154 deletions
--- a/compiler/include/concretelang/Conversion/TFHEToConcrete/Patterns.h
+++ b/compiler/include/concretelang/Conversion/TFHEToConcrete/Patterns.h
@@ -25,8 +25,8 @@ LweCiphertextType convertTypeToLWE(mlir::MLIRContext *context,
                                   mlir::Type type) {
  auto glwe = type.dyn_cast_or_null<GLWECipherTextType>();
  if (glwe != nullptr) {
-    return LweCiphertextType::get(
-        context, glwe.getDimension() * glwe.getPolynomialSize(), glwe.getP());
+    assert(glwe.getPolynomialSize() == 1);
+    return LweCiphertextType::get(context, glwe.getDimension(), glwe.getP());
  }
  auto lwe = type.dyn_cast_or_null<LweCiphertextType>();
  if (lwe != nullptr) {
--- a/compiler/include/concretelang/Conversion/Utils/GlobalFHEContext.h
+++ b/compiler/include/concretelang/Conversion/Utils/GlobalFHEContext.h
@@ -34,9 +34,9 @@ struct V0Parameter {
        ksLevel(ksLevel), ksLogBase(ksLogBase) {}

  // TODO remove the shift when we have true polynomial size
-  size_t getNBigGlweDimension() {
-    return glweDimension * (1 << logPolynomialSize);
-  }
+  size_t getPolynomialSize() { return 1 << logPolynomialSize; }
+
+  size_t getNBigGlweDimension() { return glweDimension * getPolynomialSize(); }
 };

 struct V0FHEContext {
--- a/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteOps.td
+++ b/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteOps.td
@@ -38,8 +38,8 @@ def NegateLweBufferOp : BConcrete_Op<"negate_lwe_buffer"> {
 def FillGlweFromTable : BConcrete_Op<"fill_glwe_from_table"> {
    let arguments = (ins
        1DTensorOf<[I64]>:$glwe,
-        I32Attr:$polynomialSize,
        I32Attr:$glweDimension,
+        I32Attr:$polynomialSize,
        I32Attr:$outPrecision,
        1DTensorOf<[I64]>:$table
    );
@@ -61,8 +61,6 @@ def BootstrapLweBufferOp : BConcrete_Op<"bootstrap_lwe_buffer"> {
        // LweBootstrapKeyType:$bootstrap_key,
        1DTensorOf<[I64]>:$input_ciphertext,
        1DTensorOf<[I64]>:$accumulator,
-        I32Attr:$glweDimension,
-        I32Attr:$polynomialSize,
        I32Attr:$level,
        I32Attr:$baseLog
    );
--- a/compiler/include/concretelang/Dialect/Concrete/IR/ConcreteOps.td
+++ b/compiler/include/concretelang/Dialect/Concrete/IR/ConcreteOps.td
@@ -65,8 +65,6 @@ def BootstrapLweOp : Concrete_Op<"bootstrap_lwe"> {
    let arguments = (ins 
        LweCiphertextType:$input_ciphertext,
        GlweCiphertextType:$accumulator,
-        I32Attr:$glweDimension,
-        I32Attr:$polynomialSize,
        I32Attr:$level,
        I32Attr:$baseLog
    );
--- a/compiler/include/concretelang/Dialect/Concrete/IR/ConcreteTypes.td
+++ b/compiler/include/concretelang/Dialect/Concrete/IR/ConcreteTypes.td
@@ -5,85 +5,89 @@ include "mlir/IR/BuiltinTypes.td"

 include "concretelang/Dialect/Concrete/IR/ConcreteDialect.td"

-class Concrete_Type<string name, list<Trait> traits = []>
-    : TypeDef<Concrete_Dialect, name, traits> {}
+class Concrete_Type<string name, list<Trait> traits = []> : TypeDef<Concrete_Dialect, name, traits> { }

 def GlweCiphertextType : Concrete_Type<"GlweCiphertext"> {
-  let mnemonic = "glwe_ciphertext";
+    let mnemonic = "glwe_ciphertext";

-  let summary = "A GLWE ciphertext (encryption of a polynomial of "
-                "fixed-precision integers)";
+    let summary = "A GLWE ciphertext (encryption of a polynomial of fixed-precision integers)";

-  let description = [{GLWE ciphertext.}];
+    let description = [{
+       GLWE ciphertext.
+    }];

-  let hasCustomAssemblyFormat = 1;
+    let hasCustomAssemblyFormat = 1;

-  let parameters = (ins "signed"
-                    : $polynomialSize, "signed"
-                    : $glweDimension,
-                      // Precision of the lwe ciphertext
-                      "signed"
-                    : $p);
+    let parameters = (ins
+        "signed":$glweDimension,
+        "signed":$polynomialSize,
+        // Precision of the lwe ciphertext
+        "signed":$p
+    );
 }

-def LweCiphertextType
-    : Concrete_Type<"LweCiphertext", [MemRefElementTypeInterface]> {
-  let mnemonic = "lwe_ciphertext";
+def LweCiphertextType : Concrete_Type<"LweCiphertext", [MemRefElementTypeInterface]> {
+    let mnemonic = "lwe_ciphertext";

-  let summary = "A LWE ciphertext (encryption of a fixed-precision integer)";
+    let summary = "A LWE ciphertext (encryption of a fixed-precision integer)";

-  let description = [{Learning With Error ciphertext.}];
+    let description = [{
+       Learning With Error ciphertext.
+    }];

-  let parameters = (ins
-                    // The dimension of the lwe ciphertext
-                    "signed"
-                    : $dimension,
-                      // Precision of the lwe ciphertext
-                      "signed"
-                    : $p);

-  let hasCustomAssemblyFormat = 1;
+    let parameters = (ins
+        // The dimension of the lwe ciphertext
+        "signed":$dimension,
+        // Precision of the lwe ciphertext
+        "signed":$p
+    );
+
+    let hasCustomAssemblyFormat = 1;
 }

 def CleartextType : Concrete_Type<"Cleartext"> {
-  let mnemonic = "cleartext";
+    let mnemonic = "cleartext";

-  let summary = "A cleartext (a fixed-precision integer) ready to be "
-                "multiplied to a LWE ciphertext";
+    let summary = "A cleartext (a fixed-precision integer) ready to be multiplied to a LWE ciphertext";

-  let description = [{Cleartext.}];
+    let description = [{
+       Cleartext.
+    }];

-  let parameters = (ins
-                    // Number of bits of the cleartext representation
-                    "signed"
-                    : $p);
+    let parameters = (ins
+        // Number of bits of the cleartext representation
+        "signed":$p
+    );

-  let hasCustomAssemblyFormat = 1;
+    let hasCustomAssemblyFormat = 1;
 }

 def PlaintextType : Concrete_Type<"Plaintext"> {
-  let mnemonic = "plaintext";
+    let mnemonic = "plaintext";

-  let summary = "A Plaintext (a fixed-precision integer) ready to be added to "
-                "a LWE ciphertext";
+    let summary = "A Plaintext (a fixed-precision integer) ready to be added to a LWE ciphertext";

-  let description = [{Plaintext.}];
+    let description = [{
+       Plaintext.
+    }];

-  let parameters = (ins
-                    // Number of bits of the cleartext representation
-                    "signed"
-                    : $p);
+    let parameters = (ins
+        // Number of bits of the cleartext representation
+        "signed":$p
+    );

-  let hasCustomAssemblyFormat = 1;
+    let hasCustomAssemblyFormat = 1;
 }

 def Context : Concrete_Type<"Context"> {
-  let mnemonic = "context";
+    let mnemonic = "context";

-  let summary = "A runtime context";
+    let summary = "A runtime context";

-  let description = [{An abstract runtime context to pass contextual value,
-                      like public keys, ...}];
+    let description = [{
+       An abstract runtime context to pass contextual value, like public keys, ...
+    }];
 }

 #endif
--- a/compiler/include/concretelang/Dialect/TFHE/IR/TFHEOps.td
+++ b/compiler/include/concretelang/Dialect/TFHE/IR/TFHEOps.td
@@ -82,10 +82,12 @@ def KeySwitchGLWEOp : TFHE_Op<"keyswitch_glwe"> {
  let summary = "Change the encryption parameters of a glwe ciphertext by "
                "applying a keyswitch";

-  let arguments = (ins GLWECipherTextType
-                   : $ciphertext, I32Attr
-                   : $level, I32Attr
-                   : $baseLog);
+  let arguments = (ins 
+    GLWECipherTextType : $ciphertext,
+    I32Attr : $level,
+    I32Attr : $baseLog
+  );
+  

  let results = (outs GLWECipherTextType : $result);
 }
@@ -103,13 +105,13 @@ def BootstrapGLWEOp : TFHE_Op<"bootstrap_glwe"> {
  let summary =
      "Programmable bootstraping of a GLWE ciphertext with a lookup table";

-  let arguments = (ins GLWECipherTextType
-                   : $ciphertext, GLWECipherTextType
-                   : $lookup_table, I32Attr
-                   : $glweDimension, I32Attr
-                   : $polynomialSize, I32Attr
-                   : $level, I32Attr
-                   : $baseLog);
+  let arguments = (ins
+    GLWECipherTextType : $ciphertext,
+    GLWECipherTextType : $lookup_table,
+    I32Attr : $level,
+    I32Attr : $baseLog
+  );
+
  let results = (outs GLWECipherTextType : $result);
 }

--- a/compiler/include/concretelang/Dialect/TFHE/IR/TFHETypes.td
+++ b/compiler/include/concretelang/Dialect/TFHE/IR/TFHETypes.td
@@ -31,6 +31,16 @@ def GLWECipherTextType : TFHE_Type<"GLWECipherText", [MemRefElementTypeInterface
    let hasCustomAssemblyFormat = 1;

    let genVerifyDecl = true;
+
+    let extraClassDeclaration = [{
+        // Returns true if has an unparametrized parameters
+        bool hasUnparametrizedParameters() {
+            return getDimension() ==-1 ||
+                   getPolynomialSize() == -1 ||
+                   getBits() == -1 ||
+                   getP() == -1;
+        };
+    }];
 }

 #endif
--- a/compiler/include/concretelang/Support/CompilerEngine.h
+++ b/compiler/include/concretelang/Support/CompilerEngine.h
@@ -40,6 +40,8 @@ protected:
 struct CompilationOptions {
  llvm::Optional<mlir::concretelang::V0FHEConstraint> v0FHEConstraints;

+  llvm::Optional<mlir::concretelang::V0Parameter> v0Parameter;
+
  bool verifyDiagnostics;

  bool autoParallelize;
--- a/compiler/lib/Conversion/ConcreteToBConcrete/ConcreteToBConcrete.cpp
+++ b/compiler/lib/Conversion/ConcreteToBConcrete/ConcreteToBConcrete.cpp
@@ -299,7 +299,7 @@ struct GlweFromTablePattern : public mlir::OpRewritePattern<
    auto outPrecision = resultTy.getP();

    rewriter.create<mlir::concretelang::BConcrete::FillGlweFromTable>(
-        op.getLoc(), init, polySize, glweDimension, outPrecision, op.table());
+        op.getLoc(), init, glweDimension, polySize, outPrecision, op.table());

    return ::mlir::success();
  };
--- a/compiler/lib/Conversion/FHEToTFHE/FHEToTFHE.cpp
+++ b/compiler/lib/Conversion/FHEToTFHE/FHEToTFHE.cpp
@@ -110,7 +110,7 @@ struct ApplyLookupTableEintOpPattern
        });
    //  %0 = "TFHE.bootstrap_glwe"(%glwe_ks, %glwe_lut)
    rewriter.replaceOpWithNewOp<TFHE::BootstrapGLWEOp>(lutOp, resultTy, glweKs,
-                                                       glweLut, -1, -1, -1, -1);
+                                                       glweLut, -1, -1);
    return ::mlir::success();
  };
 };
--- a/compiler/lib/Conversion/TFHEGlobalParametrization/TFHEGlobalParametrization.cpp
+++ b/compiler/lib/Conversion/TFHEGlobalParametrization/TFHEGlobalParametrization.cpp
@@ -37,20 +37,17 @@ class TFHEGlobalParametrizationTypeConverter : public mlir::TypeConverter {

 public:
  TFHEGlobalParametrizationTypeConverter(
-      mlir::concretelang::V0FHEContext &fheContext) {
+      mlir::concretelang::V0FHEContext &fheContext)
+      : fheContext(fheContext) {
    auto convertGLWECiphertextType =
-        [](GLWECipherTextType type,
-           mlir::concretelang::V0FHEContext &fheContext) {
-          auto glweDimension = fheContext.parameter.getNBigGlweDimension();
-          auto p = fheContext.constraint.p;
-          if (type.getDimension() == (signed)glweDimension &&
-              type.getP() == (signed)p) {
+        [&](GLWECipherTextType type,
+            mlir::concretelang::V0FHEContext &fheContext) {
+          auto newTy = this->glweInterPBSType(type.getContext(), fheContext);
+          if (newTy.getDimension() == type.getDimension() &&
+              newTy.getPolynomialSize() == type.getPolynomialSize() &&
+              newTy.getP() == type.getP())
            return type;
-          }
-          return GLWECipherTextType::get(
-              type.getContext(), glweDimension,
-              1 /*for the v0, is always lwe ciphertext*/,
-              64 /*for the v0 we handle only q=64*/, p);
+          return newTy;
        };
    addConversion([](mlir::Type type) { return type; });
    addConversion([&](GLWECipherTextType type) {
@@ -66,12 +63,33 @@ public:
      return r;
    });
  }
+
+  TFHE::GLWECipherTextType
+  glweInterPBSType(mlir::MLIRContext *context,
+                   mlir::concretelang::V0FHEContext fheContext) {
+    return TFHE::GLWECipherTextType::get(
+        context, fheContext.parameter.getNBigGlweDimension(), 1, 64,
+        fheContext.constraint.p);
+  }
+
+  TFHE::GLWECipherTextType glweLookupTableType(mlir::MLIRContext *context) {
+    return TFHE::GLWECipherTextType::get(
+        context, fheContext.parameter.glweDimension,
+        fheContext.parameter.getPolynomialSize(), 64, fheContext.constraint.p);
+  }
+
+  TFHE::GLWECipherTextType glweIntraPBSType(mlir::MLIRContext *context) {
+    return TFHE::GLWECipherTextType::get(context, fheContext.parameter.nSmall,
+                                         1, 64, fheContext.constraint.p);
+  }
+
+  mlir::concretelang::V0FHEContext fheContext;
 };

 struct KeySwitchGLWEOpPattern
    : public mlir::OpRewritePattern<TFHE::KeySwitchGLWEOp> {
  KeySwitchGLWEOpPattern(mlir::MLIRContext *context,
-                         mlir::TypeConverter &converter,
+                         TFHEGlobalParametrizationTypeConverter &converter,
                         mlir::concretelang::V0FHEContext &fheContext,
                         mlir::PatternBenefit benefit =
                             mlir::concretelang::DEFAULT_PATTERN_BENEFIT)
@@ -83,9 +101,7 @@ struct KeySwitchGLWEOpPattern
                  mlir::PatternRewriter &rewriter) const override {
    mlir::SmallVector<mlir::Type, 1> newResultTypes;
    auto inputTy = ksOp.ciphertext().getType().cast<TFHE::GLWECipherTextType>();
-    auto outputTy = rewriter.getType<TFHE::GLWECipherTextType>(
-        fheContext.parameter.glweDimension, fheContext.parameter.nSmall, 64,
-        fheContext.constraint.p);
+    auto outputTy = converter.glweIntraPBSType(rewriter.getContext());
    auto newOp = rewriter.replaceOpWithNewOp<TFHE::KeySwitchGLWEOp>(
        ksOp, outputTy, ksOp.ciphertext(), fheContext.parameter.ksLevel,
        fheContext.parameter.ksLogBase);
@@ -96,14 +112,14 @@ struct KeySwitchGLWEOpPattern
  };

 private:
-  mlir::TypeConverter &converter;
+  TFHEGlobalParametrizationTypeConverter &converter;
  mlir::concretelang::V0FHEContext &fheContext;
 };

 struct BootstrapGLWEOpPattern
    : public mlir::OpRewritePattern<TFHE::BootstrapGLWEOp> {
  BootstrapGLWEOpPattern(mlir::MLIRContext *context,
-                         mlir::TypeConverter &converter,
+                         TFHEGlobalParametrizationTypeConverter &converter,
                         mlir::concretelang::V0FHEContext &fheContext,
                         mlir::PatternBenefit benefit =
                             mlir::concretelang::DEFAULT_PATTERN_BENEFIT)
@@ -115,22 +131,19 @@ struct BootstrapGLWEOpPattern
                  mlir::PatternRewriter &rewriter) const override {
    auto newOp = rewriter.replaceOpWithNewOp<TFHE::BootstrapGLWEOp>(
        bsOp, converter.convertType(bsOp.result().getType()), bsOp.ciphertext(),
-        bsOp.lookup_table(), fheContext.parameter.glweDimension,
-        1 << fheContext.parameter.logPolynomialSize,
-        fheContext.parameter.brLevel, fheContext.parameter.brLogBase);
+        bsOp.lookup_table(), fheContext.parameter.brLevel,
+        fheContext.parameter.brLogBase);
    rewriter.startRootUpdate(newOp);
-    auto newInputTy = rewriter.getType<TFHE::GLWECipherTextType>(
-        fheContext.parameter.glweDimension, fheContext.parameter.nSmall, 64,
-        fheContext.constraint.p);
-    newOp.ciphertext().setType(newInputTy);
+    newOp.ciphertext().setType(
+        converter.glweIntraPBSType(rewriter.getContext()));
    newOp.lookup_table().setType(
-        converter.convertType(newOp.lookup_table().getType()));
+        converter.glweLookupTableType(rewriter.getContext()));
    rewriter.finalizeRootUpdate(newOp);
    return mlir::success();
  };

 private:
-  mlir::TypeConverter &converter;
+  TFHEGlobalParametrizationTypeConverter &converter;
  mlir::concretelang::V0FHEContext &fheContext;
 };

@@ -156,7 +169,7 @@ private:
 struct GLWEFromTablePattern
    : public mlir::OpRewritePattern<TFHE::GLWEFromTableOp> {
  GLWEFromTablePattern(mlir::MLIRContext *context,
-                       mlir::TypeConverter &converter,
+                       TFHEGlobalParametrizationTypeConverter &converter,
                       mlir::concretelang::V0FHEContext &fheContext,
                       mlir::PatternBenefit benefit =
                           mlir::concretelang::DEFAULT_PATTERN_BENEFIT)
@@ -166,8 +179,7 @@ struct GLWEFromTablePattern
  mlir::LogicalResult
  matchAndRewrite(TFHE::GLWEFromTableOp glweOp,
                  mlir::PatternRewriter &rewriter) const override {
-    auto newTy = converter.convertType(glweOp.getType())
-                     .cast<TFHE::GLWECipherTextType>();
+    auto newTy = converter.glweLookupTableType(glweOp.getContext());

    auto lutOp = glweOp.table();
    auto tableTy = lutOp.getType().cast<mlir::RankedTensorType>();
@@ -208,7 +220,7 @@ struct GLWEFromTablePattern
  };

 private:
-  mlir::TypeConverter &converter;
+  TFHEGlobalParametrizationTypeConverter &converter;
  mlir::concretelang::V0FHEContext &fheContext;
 };

@@ -269,7 +281,9 @@ void TFHEGlobalParametrizationPass::runOnOperation() {
    patterns.add<GLWEFromTablePattern>(&getContext(), converter, fheContext);
    target.addDynamicallyLegalOp<TFHE::GLWEFromTableOp>(
        [&](TFHE::GLWEFromTableOp op) {
-          return converter.isLegal(op->getResultTypes());
+          return !op.getType()
+                      .cast<TFHE::GLWECipherTextType>()
+                      .hasUnparametrizedParameters();
        });
    target.addLegalOp<mlir::arith::ConstantOp>();
    patterns.add<KeySwitchGLWEOpPattern>(&getContext(), converter, fheContext);
--- a/compiler/lib/Conversion/TFHEToConcrete/TFHEToConcrete.cpp
+++ b/compiler/lib/Conversion/TFHEToConcrete/TFHEToConcrete.cpp
@@ -88,8 +88,8 @@ struct BootstrapGLWEOpPattern
    mlir::Type resultType = converter.convertType(bsOp.getType());

    auto newOp = rewriter.replaceOpWithNewOp<Concrete::BootstrapLweOp>(
-        bsOp, resultType, bsOp.ciphertext(), bsOp.lookup_table(), -1, -1,
-        bsOp.level(), bsOp.baseLog());
+        bsOp, resultType, bsOp.ciphertext(), bsOp.lookup_table(), bsOp.level(),
+        bsOp.baseLog());

    rewriter.startRootUpdate(newOp);

--- a/compiler/lib/Dialect/Concrete/IR/ConcreteDialect.cpp
+++ b/compiler/lib/Dialect/Concrete/IR/ConcreteDialect.cpp
@@ -28,13 +28,13 @@ void ConcreteDialect::initialize() {
 mlir::Type GlweCiphertextType::parse(mlir::AsmParser &parser) {
  if (parser.parseLess())
    return Type();
-  int polynomialSize = -1;
-  if (parser.parseOptionalKeyword("_") && parser.parseInteger(polynomialSize))
+  int glweDimension = -1;
+  if (parser.parseOptionalKeyword("_") && parser.parseInteger(glweDimension))
    return Type();
  if (parser.parseComma())
    return Type();
-  int glweDimension = -1;
-  if (parser.parseOptionalKeyword("_") && parser.parseInteger(glweDimension))
+  int polynomialSize = -1;
+  if (parser.parseOptionalKeyword("_") && parser.parseInteger(polynomialSize))
    return Type();
  if (parser.parseComma())
    return Type();
@@ -45,21 +45,21 @@ mlir::Type GlweCiphertextType::parse(mlir::AsmParser &parser) {
  if (parser.parseGreater())
    return Type();
  Location loc = parser.getEncodedSourceLoc(parser.getNameLoc());
-  return getChecked(loc, loc.getContext(), polynomialSize, glweDimension, p);
+  return getChecked(loc, loc.getContext(), glweDimension, polynomialSize, p);
 }

 void GlweCiphertextType::print(mlir::AsmPrinter &p) const {
  p << "<";
-  if (getImpl()->polynomialSize == -1)
-    p << "_";
-  else
-    p << getImpl()->polynomialSize;
-  p << ",";
  if (getImpl()->glweDimension == -1)
    p << "_";
  else
    p << getImpl()->glweDimension;
  p << ",";
+  if (getImpl()->polynomialSize == -1)
+    p << "_";
+  else
+    p << getImpl()->polynomialSize;
+  p << ",";
  if (getImpl()->p == -1)
    p << "_";
  else
--- a/compiler/lib/Support/CompilerEngine.cpp
+++ b/compiler/lib/Support/CompilerEngine.cpp
@@ -144,17 +144,22 @@ llvm::Error CompilerEngine::determineFHEParameters(CompilationResult &res) {
  if (!fheConstraintOrErr.get().hasValue()) {
    return llvm::Error::success();
  }
-  auto fheParams = getV0Parameter(fheConstraintOrErr.get().getValue(),
-                                  this->compilerOptions.optimizerConfig);
+  llvm::Optional<V0Parameter> v0Params;
+  if (compilerOptions.v0Parameter.hasValue()) {
+    v0Params = compilerOptions.v0Parameter;
+  } else {
+    v0Params = getV0Parameter(fheConstraintOrErr.get().getValue(),
+                              this->compilerOptions.optimizerConfig);

-  if (!fheParams) {
-    return StreamStringError()
-           << "Could not determine V0 parameters for 2-norm of "
-           << (*fheConstraintOrErr)->norm2 << " and p of "
-           << (*fheConstraintOrErr)->p;
+    if (!v0Params) {
+      return StreamStringError()
+             << "Could not determine V0 parameters for 2-norm of "
+             << (*fheConstraintOrErr)->norm2 << " and p of "
+             << (*fheConstraintOrErr)->p;
+    }
  }
  res.fheContext.emplace(mlir::concretelang::V0FHEContext{
-      (*fheConstraintOrErr).getValue(), fheParams.getValue()});
+      (*fheConstraintOrErr).getValue(), v0Params.getValue()});

  return llvm::Error::success();
 }
--- a/compiler/src/main.cpp
+++ b/compiler/src/main.cpp
@@ -21,6 +21,7 @@

 #include "concretelang/ClientLib/KeySet.h"
 #include "concretelang/ClientLib/KeySetCache.h"
+#include "concretelang/Common/Error.h"
 #include "concretelang/Conversion/Passes.h"
 #include "concretelang/Conversion/Utils/GlobalFHEContext.h"
 #include "concretelang/Dialect/Concrete/IR/ConcreteDialect.h"
@@ -175,16 +176,6 @@ llvm::cl::opt<std::string> jitKeySetCachePath(
    "jit-keyset-cache-path",
    llvm::cl::desc("Path to cache KeySet content (unsecure)"));

-llvm::cl::opt<llvm::Optional<size_t>, false, OptionalSizeTParser>
-    assumeMaxEintPrecision(
-        "assume-max-eint-precision",
-        llvm::cl::desc("Assume a maximum precision for encrypted integers"));
-
-llvm::cl::opt<llvm::Optional<size_t>, false, OptionalSizeTParser> assumeMaxMANP(
-    "assume-max-manp",
-    llvm::cl::desc(
-        "Assume a maximum for the Minimum Arithmetic Noise Padding"));
-
 llvm::cl::opt<double> pbsErrorProbability(
    "pbs-error-probability",
    llvm::cl::desc("Change the default probability of error for all pbs"),
@@ -200,6 +191,20 @@ llvm::cl::list<int64_t> fhelinalgTileSizes(
    llvm::cl::desc(
        "Force tiling of FHELinalg operation with the given tile sizes"),
    llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated);
+
+llvm::cl::list<size_t> v0Constraint(
+    "v0-constraint",
+    llvm::cl::desc(
+        "Force the compiler to use the given v0 constraint [p, norm2]"),
+    llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated);
+
+llvm::cl::list<int64_t> v0Parameter(
+    "v0-parameter",
+    llvm::cl::desc(
+        "Force to apply the given v0 parameters [glweDimension, "
+        "logPolynomialSize, nSmall, brLevel, brLobBase, ksLevel, ksLogBase]"),
+    llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated);
+
 } // namespace cmdline

 namespace llvm {
@@ -219,7 +224,8 @@ llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
 }
 } // namespace llvm

-mlir::concretelang::CompilationOptions cmdlineCompilationOptions() {
+llvm::Expected<mlir::concretelang::CompilationOptions>
+cmdlineCompilationOptions() {
  mlir::concretelang::CompilationOptions options;

  options.verifyDiagnostics = cmdline::verifyDiagnostics;
@@ -228,12 +234,14 @@ mlir::concretelang::CompilationOptions cmdlineCompilationOptions() {
  options.dataflowParallelize = cmdline::dataflowParallelize;
  options.optimizeConcrete = cmdline::optimizeConcrete;

-  if (cmdline::assumeMaxEintPrecision.hasValue() &&
-      cmdline::assumeMaxMANP.hasValue()) {
+  if (!cmdline::v0Constraint.empty()) {
+    if (cmdline::v0Constraint.size() != 2) {
+      return llvm::make_error<llvm::StringError>(
+          "The v0-constraint option expect a list of size 2",
+          llvm::inconvertibleErrorCode());
+    }
    options.v0FHEConstraints = mlir::concretelang::V0FHEConstraint{
-        cmdline::assumeMaxMANP.getValue().getValue(),
-        cmdline::assumeMaxEintPrecision.getValue().getValue(),
-    };
+        cmdline::v0Constraint[1], cmdline::v0Constraint[0]};
  }

  if (!cmdline::funcName.empty()) {
@@ -244,6 +252,19 @@ mlir::concretelang::CompilationOptions cmdlineCompilationOptions() {
  if (!cmdline::fhelinalgTileSizes.empty())
    options.fhelinalgTileSizes.emplace(cmdline::fhelinalgTileSizes);

+  if (!cmdline::v0Parameter.empty()) {
+    if (cmdline::v0Parameter.size() != 7) {
+      return llvm::make_error<llvm::StringError>(
+          "The v0-parameter option expect a list of size 7",
+          llvm::inconvertibleErrorCode());
+    }
+    options.v0Parameter = mlir::concretelang::V0Parameter(
+        cmdline::v0Parameter[0], cmdline::v0Parameter[1],
+        cmdline::v0Parameter[2], cmdline::v0Parameter[3],
+        cmdline::v0Parameter[4], cmdline::v0Parameter[5],
+        cmdline::v0Parameter[6]);
+  }
+
  options.optimizerConfig.p_error = cmdline::pbsErrorProbability;
  options.optimizerConfig.display = cmdline::displayOptimizerChoice;

@@ -410,6 +431,10 @@ mlir::LogicalResult compilerMain(int argc, char **argv) {
  }

  auto compilerOptions = cmdlineCompilationOptions();
+  if (auto err = compilerOptions.takeError()) {
+    llvm::errs() << err << "\n";
+    return mlir::failure();
+  }

  llvm::Optional<clientlib::KeySetCache> jitKeySetCache;
  if (!cmdline::jitKeySetCachePath.empty()) {
@@ -447,7 +472,7 @@ mlir::LogicalResult compilerMain(int argc, char **argv) {
    auto process = [&](std::unique_ptr<llvm::MemoryBuffer> inputBuffer,
                       llvm::raw_ostream &os) {
      return processInputBuffer(
-          std::move(inputBuffer), fileName, compilerOptions, cmdline::action,
+          std::move(inputBuffer), fileName, *compilerOptions, cmdline::action,
          cmdline::jitArgs, jitKeySetCache, os, outputLib);
    };
    auto &os = output->os();
--- a/compiler/tests/Conversion/ConcreteToBConcrete/apply_lookup_table.mlir
+++ b/compiler/tests/Conversion/ConcreteToBConcrete/apply_lookup_table.mlir
@@ -8,8 +8,8 @@
 //CHECK:   return %[[V2]] : tensor<1025xi64>
 //CHECK: }
 func @apply_lookup_table(%arg0: !Concrete.lwe_ciphertext<1024,4>, %arg1: tensor<16xi64>) -> !Concrete.lwe_ciphertext<1024,4> {
-  %0 = "Concrete.glwe_from_table"(%arg1) {glweDimension = 1 : i32, p = 4 : i32, polynomialSize = 1024 : i32} : (tensor<16xi64>) -> !Concrete.glwe_ciphertext<1024,1,4>
+  %0 = "Concrete.glwe_from_table"(%arg1) {glweDimension = 1 : i32, p = 4 : i32, polynomialSize = 1024 : i32} : (tensor<16xi64>) -> !Concrete.glwe_ciphertext<1,1024,4>
  %1 = "Concrete.keyswitch_lwe"(%arg0) {baseLog = 2 : i32, inputLweDimension = 1 : i32, level = 3 : i32, outputLweDimension = 600 : i32} : (!Concrete.lwe_ciphertext<1024,4>) -> !Concrete.lwe_ciphertext<600,4>
-  %2 = "Concrete.bootstrap_lwe"(%1, %0) {baseLog = 4 : i32, glweDimension = 1 : i32, level = 5 : i32, polynomialSize = 1024 : i32} : (!Concrete.lwe_ciphertext<600,4>, !Concrete.glwe_ciphertext<1024,1,4>) -> !Concrete.lwe_ciphertext<1024,4>
+  %2 = "Concrete.bootstrap_lwe"(%1, %0) {baseLog = 4 : i32, glweDimension = 1 : i32, level = 5 : i32, polynomialSize = 1024 : i32} : (!Concrete.lwe_ciphertext<600,4>, !Concrete.glwe_ciphertext<1,1024,4>) -> !Concrete.lwe_ciphertext<1024,4>
  return %2 : !Concrete.lwe_ciphertext<1024,4>
 }
--- a/compiler/tests/Conversion/ConcreteToBConcrete/apply_lookup_table_cst.mlir
+++ b/compiler/tests/Conversion/ConcreteToBConcrete/apply_lookup_table_cst.mlir
@@ -10,8 +10,8 @@
 //CHECK: }
 func @apply_lookup_table_cst(%arg0: !Concrete.lwe_ciphertext<2048,4>) -> !Concrete.lwe_ciphertext<2048,4> {
  %tlu = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : tensor<16xi64>
-  %0 = "Concrete.glwe_from_table"(%tlu) {glweDimension = 1 : i32, p = 4 : i32, polynomialSize = 2048 : i32} : (tensor<16xi64>) -> !Concrete.glwe_ciphertext<2048,1,4>
+  %0 = "Concrete.glwe_from_table"(%tlu) {glweDimension = 1 : i32, p = 4 : i32, polynomialSize = 2048 : i32} : (tensor<16xi64>) -> !Concrete.glwe_ciphertext<1, 2048,4>
  %1 = "Concrete.keyswitch_lwe"(%arg0) {baseLog = 2 : i32, inputLweDimension = 1 : i32, level = 3 : i32, outputLweDimension = 600 : i32} : (!Concrete.lwe_ciphertext<2048,4>) -> !Concrete.lwe_ciphertext<600,4>
-  %2 = "Concrete.bootstrap_lwe"(%1, %0) {baseLog = 4 : i32, glweDimension = 1 : i32, level = 5 : i32, polynomialSize = 2048 : i32} : (!Concrete.lwe_ciphertext<600,4>, !Concrete.glwe_ciphertext<2048,1,4>) -> !Concrete.lwe_ciphertext<2048,4>
+  %2 = "Concrete.bootstrap_lwe"(%1, %0) {baseLog = 4 : i32, glweDimension = 1 : i32, level = 5 : i32, polynomialSize = 2048 : i32} : (!Concrete.lwe_ciphertext<600,4>, !Concrete.glwe_ciphertext<1,2048,4>) -> !Concrete.lwe_ciphertext<2048,4>
  return %2 : !Concrete.lwe_ciphertext<2048,4>
 }
--- a/compiler/tests/Conversion/FHEToTFHE/FHEToTFHE/apply_univariate.mlir
+++ b/compiler/tests/Conversion/FHEToTFHE/FHEToTFHE/apply_univariate.mlir
@@ -3,7 +3,7 @@
 // CHECK: func @apply_lookup_table(%[[A0:.*]]: !TFHE.glwe<{_,_,_}{2}>, %[[LUT:.*]]: tensor<4xi64>) -> !TFHE.glwe<{_,_,_}{3}> {
 // CHECK-NEXT: %[[V0:.*]] = "TFHE.glwe_from_table"(%[[LUT]]) : (tensor<4xi64>) -> !TFHE.glwe<{_,_,_}{2}>
 // CHECK-NEXT: %[[V1:.*]] = "TFHE.keyswitch_glwe"(%[[A0]]) {baseLog = -1 : i32, level = -1 : i32} : (!TFHE.glwe<{_,_,_}{2}>) -> !TFHE.glwe<{_,_,_}{2}>
-// CHECK-NEXT: %[[V2:.*]] = "TFHE.bootstrap_glwe"(%[[V1]], %[[V0]]) {baseLog = -1 : i32, glweDimension = -1 : i32, level = -1 : i32, polynomialSize = -1 : i32} : (!TFHE.glwe<{_,_,_}{2}>, !TFHE.glwe<{_,_,_}{2}>) -> !TFHE.glwe<{_,_,_}{3}>
+// CHECK-NEXT: %[[V2:.*]] = "TFHE.bootstrap_glwe"(%[[V1]], %[[V0]]) {baseLog = -1 : i32, level = -1 : i32} : (!TFHE.glwe<{_,_,_}{2}>, !TFHE.glwe<{_,_,_}{2}>) -> !TFHE.glwe<{_,_,_}{3}>
 // CHECK-NEXT: return %[[V2]] : !TFHE.glwe<{_,_,_}{3}>
 func @apply_lookup_table(%arg0: !FHE.eint<2>, %arg1: tensor<4xi64>) -> !FHE.eint<3> {
  %1 = "FHE.apply_lookup_table"(%arg0, %arg1): (!FHE.eint<2>, tensor<4xi64>) -> (!FHE.eint<3>)
--- a/compiler/tests/Conversion/FHEToTFHE/FHEToTFHE/apply_univariate_cst.mlir
+++ b/compiler/tests/Conversion/FHEToTFHE/FHEToTFHE/apply_univariate_cst.mlir
@@ -4,7 +4,7 @@
 //CHECK-NEXT:  %cst = arith.constant dense<"0x00000000000000000100000000000000020000000000000003000000000000000400000000000000050000000000000006000000000000000700000000000000080000000000000009000000000000000A000000000000000B000000000000000C000000000000000D000000000000000E000000000000000F0000000000000010000000000000001100000000000000120000000000000013000000000000001400000000000000150000000000000016000000000000001700000000000000180000000000000019000000000000001A000000000000001B000000000000001C000000000000001D000000000000001E000000000000001F0000000000000020000000000000002100000000000000220000000000000023000000000000002400000000000000250000000000000026000000000000002700000000000000280000000000000029000000000000002A000000000000002B000000000000002C000000000000002D000000000000002E000000000000002F0000000000000030000000000000003100000000000000320000000000000033000000000000003400000000000000350000000000000036000000000000003700000000000000380000000000000039000000000000003A000000000000003B000000000000003C000000000000003D000000000000003E000000000000003F0000000000000040000000000000004100000000000000420000000000000043000000000000004400000000000000450000000000000046000000000000004700000000000000480000000000000049000000000000004A000000000000004B000000000000004C000000000000004D000000000000004E000000000000004F0000000000000050000000000000005100000000000000520000000000000053000000000000005400000000000000550000000000000056000000000000005700000000000000580000000000000059000000000000005A000000000000005B000000000000005C000000000000005D000000000000005E000000000000005F0000000000000060000000000000006100000000000000620000000000000063000000000000006400000000000000650000000000000066000000000000006700000000000000680000000000000069000000000000006A000000000000006B000000000000006C000000000000006D000000000000006E000000000000006F0000000000000070000000000000007100000000000000720000000000000073000000000000007400000000000000750000000000000076000000000000007700000000000000780000000000000079000000000000007A000000000000007B000000000000007C000000000000007D000000000000007E000000000000007F00000000000000"> : tensor<128xi64>
 //CHECK-NEXT:  %[[V0:.*]] = "TFHE.glwe_from_table"(%cst) : (tensor<128xi64>) -> !TFHE.glwe<{_,_,_}{7}>
 //CHECK-NEXT:  %[[V1:.*]] = "TFHE.keyswitch_glwe"(%[[A0]]) {baseLog = -1 : i32, level = -1 : i32} : (!TFHE.glwe<{_,_,_}{7}>) -> !TFHE.glwe<{_,_,_}{7}>
-//CHECK-NEXT:  %[[V2:.*]] = "TFHE.bootstrap_glwe"(%[[V1]], %[[V0]]) {baseLog = -1 : i32, glweDimension = -1 : i32, level = -1 : i32, polynomialSize = -1 : i32} : (!TFHE.glwe<{_,_,_}{7}>, !TFHE.glwe<{_,_,_}{7}>) -> !TFHE.glwe<{_,_,_}{7}>
+//CHECK-NEXT:  %[[V2:.*]] = "TFHE.bootstrap_glwe"(%[[V1]], %[[V0]]) {baseLog = -1 : i32, level = -1 : i32} : (!TFHE.glwe<{_,_,_}{7}>, !TFHE.glwe<{_,_,_}{7}>) -> !TFHE.glwe<{_,_,_}{7}>
 //CHECK-NEXT:  return %[[V2]] : !TFHE.glwe<{_,_,_}{7}>
 //CHECK-NEXT: }
 func @apply_lookup_table_cst(%arg0: !FHE.eint<7>) -> !FHE.eint<7> {
--- a/compiler/tests/Conversion/TFHEGlobalParametrization/pbs_ks_bs.mlir
+++ b/compiler/tests/Conversion/TFHEGlobalParametrization/pbs_ks_bs.mlir
@@ -0,0 +1,16 @@
+// RUN: concretecompiler --passes tfhe-global-parametrization --action=dump-std --v0-parameter=2,10,750,1,23,3,4 -v0-constraint=4,0 %s 2>&1| FileCheck %s
+
+//CHECK: func @main(%[[A0:.*]]: !TFHE.glwe<{2048,1,64}{4}>) -> !TFHE.glwe<{2048,1,64}{4}> {
+//CHECK:   %cst = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : tensor<16xi64>
+//CHECK:   %[[V0:.*]] = "TFHE.glwe_from_table"(%cst) : (tensor<16xi64>) -> !TFHE.glwe<{2,1024,64}{4}>
+//CHECK:   %[[V1:.*]] = "TFHE.keyswitch_glwe"(%[[A0]]) {baseLog = 4 : i32, level = 3 : i32} : (!TFHE.glwe<{2048,1,64}{4}>) -> !TFHE.glwe<{750,1,64}{4}>
+//CHECK:   %[[V2:.*]] = "TFHE.bootstrap_glwe"(%[[V1]], %[[V0]]) {baseLog = 23 : i32, level = 1 : i32} : (!TFHE.glwe<{750,1,64}{4}>, !TFHE.glwe<{2,1024,64}{4}>) -> !TFHE.glwe<{2048,1,64}{4}>
+//CHECK:   return %[[V2]] : !TFHE.glwe<{2048,1,64}{4}>
+//CHECK: }
+func @main(%arg0: !TFHE.glwe<{_,_,_}{4}>) -> !TFHE.glwe<{_,_,_}{4}> {
+  %cst = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : tensor<16xi64>
+  %0 = "TFHE.glwe_from_table"(%cst) : (tensor<16xi64>) -> !TFHE.glwe<{_,_,_}{4}>
+  %1 = "TFHE.keyswitch_glwe"(%arg0) {baseLog = -1 : i32, level = -1 : i32} : (!TFHE.glwe<{_,_,_}{4}>) -> !TFHE.glwe<{_,_,_}{4}>
+  %2 = "TFHE.bootstrap_glwe"(%1, %0) {baseLog = -1 : i32, level = -1 : i32} : (!TFHE.glwe<{_,_,_}{4}>, !TFHE.glwe<{_,_,_}{4}>) -> !TFHE.glwe<{_,_,_}{4}>
+  return %2 : !TFHE.glwe<{_,_,_}{4}>
+}
--- a/compiler/tests/Conversion/TFHEToConcrete/TFHEToConcrete/bootstrap.mlir
+++ b/compiler/tests/Conversion/TFHEToConcrete/TFHEToConcrete/bootstrap.mlir
@@ -1,14 +1,14 @@
 // RUN: concretecompiler --passes tfhe-to-concrete --action=dump-concrete %s 2>&1| FileCheck %s

-//CHECK: func @bootstrap_lwe(%[[A0:.*]]: !Concrete.lwe_ciphertext<1024,7>) -> !Concrete.lwe_ciphertext<1024,4> {
+//CHECK: func @bootstrap_lwe(%[[A0:.*]]: !Concrete.lwe_ciphertext<600,7>) -> !Concrete.lwe_ciphertext<1024,4> {
 //CHECK:   %cst = arith.constant dense<"0x00000000000000000100000000000000020000000000000003000000000000000400000000000000050000000000000006000000000000000700000000000000080000000000000009000000000000000A000000000000000B000000000000000C000000000000000D000000000000000E000000000000000F0000000000000010000000000000001100000000000000120000000000000013000000000000001400000000000000150000000000000016000000000000001700000000000000180000000000000019000000000000001A000000000000001B000000000000001C000000000000001D000000000000001E000000000000001F0000000000000020000000000000002100000000000000220000000000000023000000000000002400000000000000250000000000000026000000000000002700000000000000280000000000000029000000000000002A000000000000002B000000000000002C000000000000002D000000000000002E000000000000002F0000000000000030000000000000003100000000000000320000000000000033000000000000003400000000000000350000000000000036000000000000003700000000000000380000000000000039000000000000003A000000000000003B000000000000003C000000000000003D000000000000003E000000000000003F0000000000000040000000000000004100000000000000420000000000000043000000000000004400000000000000450000000000000046000000000000004700000000000000480000000000000049000000000000004A000000000000004B000000000000004C000000000000004D000000000000004E000000000000004F0000000000000050000000000000005100000000000000520000000000000053000000000000005400000000000000550000000000000056000000000000005700000000000000580000000000000059000000000000005A000000000000005B000000000000005C000000000000005D000000000000005E000000000000005F0000000000000060000000000000006100000000000000620000000000000063000000000000006400000000000000650000000000000066000000000000006700000000000000680000000000000069000000000000006A000000000000006B000000000000006C000000000000006D000000000000006E000000000000006F0000000000000070000000000000007100000000000000720000000000000073000000000000007400000000000000750000000000000076000000000000007700000000000000780000000000000079000000000000007A000000000000007B000000000000007C000000000000007D000000000000007E000000000000007F00000000000000"> : tensor<128xi64>
 //CHECK:   %[[V0:.*]] = "Concrete.glwe_from_table"(%cst) : (tensor<128xi64>) -> !Concrete.glwe_ciphertext<1,1024,7>
-//CHECK:   %[[V1:.*]] = "Concrete.bootstrap_lwe"(%[[A0]], %[[V0]]) {baseLog = 1 : i32, glweDimension = -1 : i32, level = 3 : i32, polynomialSize = -1 : i32} : (!Concrete.lwe_ciphertext<1024,7>, !Concrete.glwe_ciphertext<1,1024,7>) -> !Concrete.lwe_ciphertext<1024,4>
+//CHECK:   %[[V1:.*]] = "Concrete.bootstrap_lwe"(%[[A0]], %[[V0]]) {baseLog = 1 : i32, level = 3 : i32} : (!Concrete.lwe_ciphertext<600,7>, !Concrete.glwe_ciphertext<1,1024,7>) -> !Concrete.lwe_ciphertext<1024,4>
 //CHECK:   return %[[V1]] : !Concrete.lwe_ciphertext<1024,4>
 //CHECK: }
-func @bootstrap_lwe(%ciphertext: !TFHE.glwe<{1,1024,64}{7}>) -> !TFHE.glwe<{1,1024,64}{4}> {
+func @bootstrap_lwe(%ciphertext: !TFHE.glwe<{600,1,64}{7}>) -> !TFHE.glwe<{1024,1,64}{4}> {
  %cst = arith.constant dense<"0x00000000000000000100000000000000020000000000000003000000000000000400000000000000050000000000000006000000000000000700000000000000080000000000000009000000000000000A000000000000000B000000000000000C000000000000000D000000000000000E000000000000000F0000000000000010000000000000001100000000000000120000000000000013000000000000001400000000000000150000000000000016000000000000001700000000000000180000000000000019000000000000001A000000000000001B000000000000001C000000000000001D000000000000001E000000000000001F0000000000000020000000000000002100000000000000220000000000000023000000000000002400000000000000250000000000000026000000000000002700000000000000280000000000000029000000000000002A000000000000002B000000000000002C000000000000002D000000000000002E000000000000002F0000000000000030000000000000003100000000000000320000000000000033000000000000003400000000000000350000000000000036000000000000003700000000000000380000000000000039000000000000003A000000000000003B000000000000003C000000000000003D000000000000003E000000000000003F0000000000000040000000000000004100000000000000420000000000000043000000000000004400000000000000450000000000000046000000000000004700000000000000480000000000000049000000000000004A000000000000004B000000000000004C000000000000004D000000000000004E000000000000004F0000000000000050000000000000005100000000000000520000000000000053000000000000005400000000000000550000000000000056000000000000005700000000000000580000000000000059000000000000005A000000000000005B000000000000005C000000000000005D000000000000005E000000000000005F0000000000000060000000000000006100000000000000620000000000000063000000000000006400000000000000650000000000000066000000000000006700000000000000680000000000000069000000000000006A000000000000006B000000000000006C000000000000006D000000000000006E000000000000006F0000000000000070000000000000007100000000000000720000000000000073000000000000007400000000000000750000000000000076000000000000007700000000000000780000000000000079000000000000007A000000000000007B000000000000007C000000000000007D000000000000007E000000000000007F00000000000000"> : tensor<128xi64>
  %glwe_lut = "TFHE.glwe_from_table"(%cst) : (tensor<128xi64>) -> !TFHE.glwe<{1,1024,64}{7}>
-  %bootstraped = "TFHE.bootstrap_glwe"(%ciphertext, %glwe_lut) {baseLog = 1 : i32, glweDimension = 1 : i32, level = 3 : i32, polynomialSize = 1024 : i32} : (!TFHE.glwe<{1,1024,64}{7}>, !TFHE.glwe<{1,1024,64}{7}>) -> !TFHE.glwe<{1,1024,64}{4}>
-  return %bootstraped : !TFHE.glwe<{1,1024,64}{4}>
+  %bootstraped = "TFHE.bootstrap_glwe"(%ciphertext, %glwe_lut) {baseLog = 1 : i32, glweDimension = 1 : i32, level = 3 : i32, polynomialSize = 1024 : i32} : (!TFHE.glwe<{600,1,64}{7}>, !TFHE.glwe<{1,1024,64}{7}>) -> !TFHE.glwe<{1024,1,64}{4}>
+  return %bootstraped : !TFHE.glwe<{1024,1,64}{4}>
 }
--- a/compiler/tests/Conversion/TFHEToConcrete/TFHEToConcrete/keyswitch.mlir
+++ b/compiler/tests/Conversion/TFHEToConcrete/TFHEToConcrete/keyswitch.mlir
@@ -4,7 +4,7 @@
 // CHECK-NEXT:   %[[V0:.*]] = "Concrete.keyswitch_lwe"(%[[A0]]) {baseLog = 3 : i32, level = 2 : i32} : (!Concrete.lwe_ciphertext<1024,2>) -> !Concrete.lwe_ciphertext<567,2>
 // CHECK-NEXT:   return %[[V0]] : !Concrete.lwe_ciphertext<567,2>
 // CHECK-NEXT: }
-func @keyswitch_glwe(%arg0: !TFHE.glwe<{1,1024,64}{2}>) -> !TFHE.glwe<{1,567,64}{2}> {
-  %0 = "TFHE.keyswitch_glwe"(%arg0) {baseLog = 3 : i32, level = 2 : i32} : (!TFHE.glwe<{1,1024,64}{2}>) -> !TFHE.glwe<{1,567,64}{2}>
-  return %0 : !TFHE.glwe<{1,567,64}{2}>
+func @keyswitch_glwe(%arg0: !TFHE.glwe<{1024,1,64}{2}>) -> !TFHE.glwe<{567,1,64}{2}> {
+  %0 = "TFHE.keyswitch_glwe"(%arg0) {baseLog = 3 : i32, level = 2 : i32} : (!TFHE.glwe<{1024,1,64}{2}>) -> !TFHE.glwe<{567,1,64}{2}>
+  return %0 : !TFHE.glwe<{567,1,64}{2}>
 }