Merge pull request #1339 from ShirChoiEX:OpenfheInplaceOperations

PiperOrigin-RevId: 724369141
google · Feb 7, 2025 · d372bf7 · d372bf7
2 parents bd8d72e + e44d36b
commit d372bf7
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 0 deletions.
diff --git a/lib/Dialect/Openfhe/IR/OpenfheOps.td b/lib/Dialect/Openfhe/IR/OpenfheOps.td
@@ -58,6 +58,21 @@ class Openfhe_BinaryOp<string mnemonic, list<Trait> traits = []>
   let results = (outs NewLWECiphertext:$output);
 }
 
+class Openfhe_BinaryInPlaceOp<string mnemonic, list<Trait> traits = []>
+  : Openfhe_Op<mnemonic, traits # [
+    AllTypesMatch<["lhs", "rhs"]>,
+    ]> {
+
+  let summary = "In-place binary operation for OpenFHE";
+
+  let arguments = (ins
+    Openfhe_CryptoContext:$cryptoContext,
+    NewLWECiphertext:$lhs,
+    NewLWECiphertext:$rhs
+  );
+}
+
+
 def GenParamsOp : Openfhe_Op<"gen_params"> {
   let description = [{
     Generates the parameters for the OpenFHE scheme.
@@ -165,6 +180,17 @@ def SubOp : Openfhe_BinaryOp<"sub",
   let summary = "OpenFHE sub operation of two ciphertexts.";
 }
 
+// In-Place Addition
+def AddInPlaceOp : Openfhe_BinaryInPlaceOp<"add_inplace"> {
+  let summary = "Performs in-place homomorphic addition,  modifying lhs.";
+}
+
+// In-Place Subtraction
+def SubInPlaceOp : Openfhe_BinaryInPlaceOp<"sub_inplace"> {
+  let summary = "Performs in-place homomorphic subtraction, modifying lhs.";
+}
+
+
 def AddPlainOp : Openfhe_Op<"add_plain",[
     Pure,
     AllTypesMatch<["ciphertext", "output"]>

diff --git a/tests/Dialect/Openfhe/IR/ops.mlir b/tests/Dialect/Openfhe/IR/ops.mlir
@@ -70,6 +70,20 @@ module {
     %out = openfhe.add %cc, %c1, %c2: (!cc, !ct, !ct) -> !ct
     return
   }
+  // CHECK-LABEL: func @test_inplace_add
+  func.func @test_inplace_add(%cc: !cc, %pt : !pt, %pk : !pk) {
+    %c1 = openfhe.encrypt %cc, %pt, %pk : (!cc, !pt, !pk) -> !ct
+    %c2 = openfhe.encrypt %cc, %pt, %pk : (!cc, !pt, !pk) -> !ct
+    openfhe.add_inplace %cc, %c1, %c2: (!cc, !ct, !ct) -> ()
+    return
+  }
+  // CHECK-LABEL: func @test_inplace_sub
+  func.func @test_inplace_sub(%cc: !cc, %pt : !pt, %pk : !pk) {
+    %c1 = openfhe.encrypt %cc, %pt, %pk : (!cc, !pt, !pk) -> !ct
+    %c2 = openfhe.encrypt %cc, %pt, %pk : (!cc, !pt, !pk) -> !ct
+    openfhe.sub_inplace %cc, %c1, %c2: (!cc, !ct, !ct) -> ()
+    return
+  }
 
   // CHECK-LABEL: func @test_sub
   func.func @test_sub(%cc : !cc, %pt : !pt, %pk: !pk) {