From 0e59f66a645f435d21d98082b3fcc32bc85cb807 Mon Sep 17 00:00:00 2001
From: Quentin Bourgerie <bourgerie.quentin@gmail.com>
Date: Fri, 7 Apr 2023 10:23:13 +0200
Subject: [PATCH] fix(compiler): FHE.to_bool and FHE.from_bool is a no op and
 not reduce the noise to the base

---
 .../lib/Dialect/FHE/Analysis/MANP.cpp         | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/compilers/concrete-compiler/compiler/lib/Dialect/FHE/Analysis/MANP.cpp b/compilers/concrete-compiler/compiler/lib/Dialect/FHE/Analysis/MANP.cpp
index fb8ffeeef..ff24a65fa 100644
--- a/compilers/concrete-compiler/compiler/lib/Dialect/FHE/Analysis/MANP.cpp
+++ b/compilers/concrete-compiler/compiler/lib/Dialect/FHE/Analysis/MANP.cpp
@@ -294,6 +294,20 @@ static llvm::APInt conservativeIntNorm2Sq(mlir::Type t) {
   return APIntWidthExtendUnsignedSq(maxVal);
 }
 
+static llvm::APInt
+getNoOpSqMANP(llvm::ArrayRef<const MANPLattice *> operandMANPs) {
+  // Come from block arg as example
+  if (operandMANPs.size() == 0) {
+    return llvm::APInt{1, 1, false};
+  }
+  assert(operandMANPs[0]->getValue().getMANP().has_value() &&
+         "Missing squared Minimal Arithmetic Noise Padding for encrypted "
+         "operands");
+
+  llvm::APInt eNorm = operandMANPs[0]->getValue().getMANP().value();
+  return eNorm;
+}
+
 /// Calculates the squared Minimal Arithmetic Noise Padding of an
 /// `FHELinalg.dot_eint_int` operation.
 static llvm::APInt getSqMANP(mlir::concretelang::FHELinalg::Dot op,
@@ -1348,11 +1362,12 @@ public:
                    llvm::dyn_cast<mlir::concretelang::FHE::MaxEintOp>(op)) {
       norm2SqEquiv = getSqMANP(maxEintOp, operands);
     } else if (llvm::isa<mlir::concretelang::FHE::ZeroEintOp>(op) ||
-               llvm::isa<mlir::concretelang::FHE::ToBoolOp>(op) ||
-               llvm::isa<mlir::concretelang::FHE::FromBoolOp>(op) ||
                llvm::isa<mlir::concretelang::FHE::ZeroTensorOp>(op) ||
                llvm::isa<mlir::concretelang::FHE::ApplyLookupTableEintOp>(op)) {
       norm2SqEquiv = llvm::APInt{1, 1, false};
+    } else if (llvm::isa<mlir::concretelang::FHE::ToBoolOp>(op) ||
+               llvm::isa<mlir::concretelang::FHE::FromBoolOp>(op)) {
+      norm2SqEquiv = getNoOpSqMANP(operands);
     }
     // FHELinalg Operators
     else if (auto dotOp =