Lattigo: Add Negate op and e2e CMUX example

Author: ZenithalHourlyRate

lib/Dialect/LWE/Conversions/LWEToLattigo/LWEToLattigo.cpp

@@ -290,7 +290,8 @@ struct ConvertRlweCommutativePlainOp : public OpConversionPattern<PlainOp> {
 };
 
 // Lattigo API enforces ciphertext, plaintext ordering.
-template <typename EvaluatorType, typename PlainOp, typename LattigoPlainOp>
+template <typename EvaluatorType, typename PlainOp, typename LattigoPlainOp,
+          typename LattigoAddOp, typename LattigoNegateOp>
 struct ConvertRlweSubPlainOp : public OpConversionPattern<PlainOp> {
   using OpConversionPattern<PlainOp>::OpConversionPattern;
 
@@ -302,21 +303,28 @@ struct ConvertRlweSubPlainOp : public OpConversionPattern<PlainOp> {
     if (failed(result)) return result;
 
     Value evaluator = result.value();
-    Value ciphertext = adaptor.getLhs();
-    Value plaintext = adaptor.getRhs();
     if (isa<lattigo::RLWECiphertextType>(adaptor.getLhs().getType())) {
       // Lattigo API enforces ciphertext, plaintext ordering, so we can use
       // LattigoPlainOp directly.
+      Value ciphertext = adaptor.getLhs();
+      Value plaintext = adaptor.getRhs();
       rewriter.replaceOpWithNewOp<LattigoPlainOp>(
           op, this->typeConverter->convertType(op.getOutput().getType()),
           evaluator, ciphertext, plaintext);
       return success();
     }
 
-    // TODO(#1623): Support this case by lowering it to Add(Negate(ciphertext),
-    // plaintext).
-    return op.emitOpError() << "subplain op does not support plaintext, "
-                               "ciphertext operand ordering";
+    // handle plaintext - ciphertext using (-ciphertext) + plaintext
+    Value plaintext = adaptor.getLhs();
+    Value ciphertext = adaptor.getRhs();
+
+    auto negated = rewriter.create<LattigoNegateOp>(
+        op.getLoc(), this->typeConverter->convertType(op.getOutput().getType()),
+        evaluator, ciphertext);
+    rewriter.replaceOpWithNewOp<LattigoAddOp>(
+        op, this->typeConverter->convertType(op.getOutput().getType()),
+        evaluator, negated, plaintext);
+    return success();
   }
 };
 
@@ -478,8 +486,10 @@ struct ConvertLWEReinterpretApplicationData
   LogicalResult matchAndRewrite(
       lwe::ReinterpretApplicationDataOp op, OpAdaptor adaptor,
       ConversionPatternRewriter &rewriter) const override {
-    // erase reinterpret underlying
-    rewriter.replaceOp(op, adaptor.getOperands()[0].getDefiningOp());
+    // Erase reinterpret application data.
+    // If operand has no defining op, we can not replace it with defining op.
+    rewriter.replaceAllOpUsesWith(op, adaptor.getOperands()[0]);
+    rewriter.eraseOp(op);
     return success();
   }
 };
@@ -498,7 +508,8 @@ using ConvertBGVAddPlainOp =
                                   lattigo::BGVAddNewOp>;
 using ConvertBGVSubPlainOp =
     ConvertRlweSubPlainOp<lattigo::BGVEvaluatorType, bgv::SubPlainOp,
-                          lattigo::BGVSubNewOp>;
+                          lattigo::BGVSubNewOp, lattigo::BGVAddNewOp,
+                          lattigo::RLWENegateNewOp>;
 using ConvertBGVMulPlainOp =
     ConvertRlweCommutativePlainOp<lattigo::BGVEvaluatorType, bgv::MulPlainOp,
                                   lattigo::BGVMulNewOp>;
@@ -545,7 +556,8 @@ using ConvertCKKSAddPlainOp =
                                   lattigo::CKKSAddNewOp>;
 using ConvertCKKSSubPlainOp =
     ConvertRlweSubPlainOp<lattigo::CKKSEvaluatorType, ckks::SubPlainOp,
-                          lattigo::CKKSSubNewOp>;
+                          lattigo::CKKSSubNewOp, lattigo::CKKSAddNewOp,
+                          lattigo::RLWENegateNewOp>;
 using ConvertCKKSMulPlainOp =
     ConvertRlweCommutativePlainOp<lattigo::CKKSEvaluatorType, ckks::MulPlainOp,
                                   lattigo::CKKSMulNewOp>;

lib/Dialect/Lattigo/IR/LattigoRLWEOps.td

@@ -151,4 +151,31 @@ def Lattigo_RLWEDropLevelOp : Lattigo_RLWEOp<"drop_level", [InplaceOpInterface]>
   let extraClassDeclaration = "int getInplaceOperandIndex() { return 2; }";
 }
 
+def Lattigo_RLWENegateNewOp : Lattigo_RLWEOp<"negate_new"> {
+  let summary = "Negate a ciphertext";
+  let arguments = (ins
+    Lattigo_RLWEEvaluator:$evaluator,
+    Lattigo_RLWECiphertext:$input
+  );
+  let results = (outs Lattigo_RLWECiphertext:$output);
+}
+
+def Lattigo_RLWENegateOp : Lattigo_RLWEOp<"negate", [InplaceOpInterface]> {
+  let summary = "Negate of a ciphertext";
+  let description = [{
+    This operation negates a ciphertext
+
+    The result will be written to the `inplace` operand. The `output` result is
+    a transitive reference to the `inplace` operand for sake of the MLIR SSA form.
+  }];
+  let arguments = (ins
+    Lattigo_RLWEEvaluator:$evaluator,
+    Lattigo_RLWECiphertext:$input,
+    Lattigo_RLWECiphertext:$inplace
+  );
+  let results = (outs Lattigo_RLWECiphertext:$output);
+
+  let extraClassDeclaration = "int getInplaceOperandIndex() { return 2; }";
+}
+
 #endif  // LIB_DIALECT_LATTIGO_IR_LATTIGORLWEOPS_TD_

lib/Dialect/Lattigo/Transforms/AllocToInplace.cpp

@@ -331,6 +331,7 @@ struct AllocToInplace : impl::AllocToInplaceBase<AllocToInplace> {
         ConvertUnaryOp<lattigo::CKKSRescaleNewOp, lattigo::CKKSRescaleOp>,
         ConvertRotateOp<lattigo::CKKSRotateNewOp, lattigo::CKKSRotateOp>,
         // RLWE
+        ConvertUnaryOp<lattigo::RLWENegateNewOp, lattigo::RLWENegateOp>,
         ConvertDropLevelOp<lattigo::RLWEDropLevelNewOp,
                            lattigo::RLWEDropLevelOp>>(context, &liveness,
                                                       &blockToStorageInfo);