add hasDeterministicResultNoise to interface

google · Nov 29, 2023 · 177798c · 177798c
1 parent ec0c9a4
commit 177798c
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 1 deletion.
diff --git a/include/Dialect/LWE/IR/LWEOps.td b/include/Dialect/LWE/IR/LWEOps.td
@@ -41,7 +41,10 @@ def LWE_EncodeOp : LWE_Op<"encode", [Pure]> {
   let hasVerifier = 1;
 }
 
-def LWE_TrivialEncryptOp: LWE_Op<"trivial_encrypt", [Pure]> {
+def LWE_TrivialEncryptOp: LWE_Op<"trivial_encrypt", [
+  Pure,
+  DeclareOpInterfaceMethods<NoisePropagationInterface>
+]> {
   let summary = "Create a trivial encryption of a plaintext.";
 
   let arguments = (ins

diff --git a/include/Interfaces/NoiseInterfaces.td b/include/Interfaces/NoiseInterfaces.td
@@ -34,6 +34,14 @@ def NoisePropagationInterface : OpInterface<"NoisePropagationInterface"> {
     "void", "inferResultNoise", (ins
       "::llvm::ArrayRef<Variance>":$argNoises,
       "::mlir::heir::SetNoiseFn":$setValueNoise)
+
+  >,
+  InterfaceMethod<[{
+      Returns true if the noise in the result op is independent of the noise in
+      its inputs. This is suitable for ops like bootstrap and initial
+      encryption.
+    }],
+    "bool", "hasDeterministicResultNoise", (ins)
   >];
 }
 

diff --git a/lib/Dialect/CGGI/IR/CGGIOps.cpp b/lib/Dialect/CGGI/IR/CGGIOps.cpp
@@ -132,6 +132,13 @@ void NotOp::inferResultNoise(llvm::ArrayRef<Variance> argNoises,
   setValueNoise(getInput(), argNoises[0]);
 }
 
+bool AndOp::hasDeterministicResultNoise() { return true; }
+bool OrOp::hasDeterministicResultNoise() { return true; }
+bool XorOp::hasDeterministicResultNoise() { return true; }
+bool Lut3Op::hasDeterministicResultNoise() { return true; }
+bool Lut2Op::hasDeterministicResultNoise() { return true; }
+bool NotOp::hasDeterministicResultNoise() { return false; }
+
 }  // namespace cggi
 }  // namespace heir
 }  // namespace mlir

diff --git a/lib/Dialect/LWE/IR/LWEDialect.cpp b/lib/Dialect/LWE/IR/LWEDialect.cpp
@@ -185,6 +185,14 @@ void AddOp::inferResultNoise(llvm::ArrayRef<Variance> argNoises,
   return setValueNoise(getResult(), result);
 }
 
+void TrivialEncryptOp::inferResultNoise(llvm::ArrayRef<Variance> argNoises,
+                                        SetNoiseFn setValueNoise) {
+  return setValueNoise(getResult(), Variance(0));
+}
+
+bool AddOp::hasDeterministicResultNoise() { return false; }
+bool TrivialEncryptOp::hasDeterministicResultNoise() { return true; }
+
 }  // namespace lwe
 }  // namespace heir
 }  // namespace mlir