fix(optimizer): fix leveled noise propagation

zama-ai · Jul 22, 2024 · af7d986 · af7d986
1 parent aadc3a6
commit af7d986
Show file tree

Hide file tree

Showing 17 changed files with 244 additions and 250 deletions.
diff --git a/compilers/concrete-compiler/compiler/include/concretelang/Dialect/FHE/IR/FHEOps.td b/compilers/concrete-compiler/compiler/include/concretelang/Dialect/FHE/IR/FHEOps.td
@@ -19,7 +19,7 @@ include "concretelang/Dialect/FHE/Interfaces/FHEInterfaces.td"
 class FHE_Op<string mnemonic, list<Trait> traits = []> :
     Op<FHE_Dialect, mnemonic, traits>;
 
-def FHE_ZeroEintOp : FHE_Op<"zero", [Pure, ConstantNoise]> {
+def FHE_ZeroEintOp : FHE_Op<"zero", [Pure, ZeroNoise]> {
     let summary = "Returns a trivial encrypted integer of 0";
 
     let description = [{
@@ -34,7 +34,7 @@ def FHE_ZeroEintOp : FHE_Op<"zero", [Pure, ConstantNoise]> {
     let results = (outs FHE_AnyEncryptedInteger:$out);
 }
 
-def FHE_ZeroTensorOp : FHE_Op<"zero_tensor", [Pure, ConstantNoise]> {
+def FHE_ZeroTensorOp : FHE_Op<"zero_tensor", [Pure, ZeroNoise]> {
     let summary = "Creates a new tensor with all elements initialized to an encrypted zero.";
 
     let description = [{
@@ -52,7 +52,7 @@ def FHE_ZeroTensorOp : FHE_Op<"zero_tensor", [Pure, ConstantNoise]> {
     let results = (outs Type<And<[TensorOf<[FHE_AnyEncryptedInteger]>.predicate, HasStaticShapePred]>>:$tensor);
 }
 
-def FHE_AddEintIntOp : FHE_Op<"add_eint_int", [Pure, BinaryEintInt, DeclareOpInterfaceMethods<Binary>]> {
+def FHE_AddEintIntOp : FHE_Op<"add_eint_int", [Pure, BinaryEintInt, AdditiveNoise,  DeclareOpInterfaceMethods<Binary>]> {
     let summary = "Adds an encrypted integer and a clear integer";
 
     let description = [{
@@ -85,7 +85,7 @@ def FHE_AddEintIntOp : FHE_Op<"add_eint_int", [Pure, BinaryEintInt, DeclareOpInt
     let hasFolder = 1;
 }
 
-def FHE_AddEintOp : FHE_Op<"add_eint", [Pure, BinaryEint, DeclareOpInterfaceMethods<BinaryEint>]> {
+def FHE_AddEintOp : FHE_Op<"add_eint", [Pure, BinaryEint, AdditiveNoise, DeclareOpInterfaceMethods<BinaryEint>]> {
     let summary = "Adds two encrypted integers";
 
     let description = [{
@@ -117,7 +117,7 @@ def FHE_AddEintOp : FHE_Op<"add_eint", [Pure, BinaryEint, DeclareOpInterfaceMeth
     let hasVerifier = 1;
 }
 
-def FHE_SubIntEintOp : FHE_Op<"sub_int_eint", [Pure, BinaryIntEint]> {
+def FHE_SubIntEintOp : FHE_Op<"sub_int_eint", [Pure, BinaryIntEint, AdditiveNoise]> {
     let summary = "Subtract an encrypted integer from a clear integer";
 
     let description = [{
@@ -149,7 +149,7 @@ def FHE_SubIntEintOp : FHE_Op<"sub_int_eint", [Pure, BinaryIntEint]> {
     let hasVerifier = 1;
 }
 
-def FHE_SubEintIntOp : FHE_Op<"sub_eint_int", [Pure, BinaryEintInt, DeclareOpInterfaceMethods<Binary>]> {
+def FHE_SubEintIntOp : FHE_Op<"sub_eint_int", [Pure, BinaryEintInt, AdditiveNoise, DeclareOpInterfaceMethods<Binary>]> {
     let summary = "Subtract a clear integer from an encrypted integer";
 
     let description = [{
@@ -182,7 +182,7 @@ def FHE_SubEintIntOp : FHE_Op<"sub_eint_int", [Pure, BinaryEintInt, DeclareOpInt
     let hasFolder = 1;
 }
 
-def FHE_SubEintOp : FHE_Op<"sub_eint", [Pure, BinaryEint, DeclareOpInterfaceMethods<BinaryEint>]> {
+def FHE_SubEintOp : FHE_Op<"sub_eint", [Pure, BinaryEint, AdditiveNoise, DeclareOpInterfaceMethods<BinaryEint>]> {
     let summary = "Subtract an encrypted integer from an encrypted integer";
 
     let description = [{
@@ -214,7 +214,7 @@ def FHE_SubEintOp : FHE_Op<"sub_eint", [Pure, BinaryEint, DeclareOpInterfaceMeth
     let hasVerifier = 1;
 }
 
-def FHE_NegEintOp : FHE_Op<"neg_eint", [Pure, UnaryEint, DeclareOpInterfaceMethods<UnaryEint>]> {
+def FHE_NegEintOp : FHE_Op<"neg_eint", [Pure, UnaryEint, AdditiveNoise, DeclareOpInterfaceMethods<UnaryEint>]> {
 
     let summary = "Negates an encrypted integer";
 

diff --git a/...s/concrete-compiler/compiler/include/concretelang/Dialect/FHE/Interfaces/FHEInterfaces.td b/...s/concrete-compiler/compiler/include/concretelang/Dialect/FHE/Interfaces/FHEInterfaces.td
@@ -37,6 +37,22 @@ def ConstantNoise : OpInterface<"ConstantNoise"> {
     let cppNamespace = "mlir::concretelang::FHE";
 }
 
+def ZeroNoise : OpInterface<"ZeroNoise"> {
+    let description = [{
+        An operation outputs a ciphertext with zero noise.
+    }];
+
+    let cppNamespace = "mlir::concretelang::FHE";
+}
+
+def AdditiveNoise : OpInterface<"AdditiveNoise"> {
+    let description = [{
+        An n-ary operation whose output noise is the unweighted sum of all input noises.
+    }];
+
+    let cppNamespace = "mlir::concretelang::FHE";
+}
+
 def UnaryEint : OpInterface<"UnaryEint"> {
     let description = [{
         A unary operation on scalars, with the operand encrypted.
@@ -63,7 +79,7 @@ def UnaryEint : OpInterface<"UnaryEint"> {
           if (auto operandTy = dyn_cast<mlir::RankedTensorType>($_op->getOpOperand(0).get().getType())) {
             return operandTy.getElementType();
           } else return $_op->getOpOperand(0).get().getType();
-        }]>        
+        }]>
     ];
 }
 
@@ -124,8 +140,8 @@ def Binary : OpInterface<"Binary"> {
           if (auto cstOp = llvm::dyn_cast_or_null<mlir::arith::ConstantOp>($_op->
                                                         getOpOperand(opNum).get().getDefiningOp()))
             return cstOp->template getAttrOfType<mlir::DenseIntElementsAttr>("value").template getValues<llvm::APInt>();
-          else return {};        
-        }]>,        
+          else return {};
+        }]>,
     ];
 }