Add in-place OpenFHE ops and in-place transform

New ops:

- `add_inplace`
- `add_plain_inplace`
- `level_reduce_inplace`
- `mod_reduce_inplace`
- `mul_const_inplace`
- `negate_inplace`
- `relin_inplace`
- `square_inplace`
- `sub_inplace`
- `sub_plain_inplace`

Note: openfhe does not appear to have in-place variants for: `EvalMult` (ciphertext or plaintext), `EvalMulNoRelin`, `EvalRotate`, `EvalAutomorphism`, `EvalBootstrap`.

New transform: `openfhe-alloc-to-inplace`

PiperOrigin-RevId: 856367329

Author: j2kun

lib/Analysis/Cpp/BUILD

@@ -0,0 +1,19 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+
+package(
+    default_applicable_licenses = ["@heir//:license"],
+    default_visibility = ["//visibility:public"],
+)
+
+cc_library(
+    name = "ConstQualifierAnalysis",
+    srcs = ["ConstQualifierAnalysis.cpp"],
+    hdrs = ["ConstQualifierAnalysis.h"],
+    deps = [
+        "@heir//lib/Utils/Tablegen:InPlaceOpInterface",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@llvm-project//mlir:Support",
+    ],
+)

lib/Analysis/Cpp/ConstQualifierAnalysis.cpp

@@ -0,0 +1,36 @@
+#include "lib/Analysis/Cpp/ConstQualifierAnalysis.h"
+
+#include "lib/Utils/Tablegen/InPlaceOpInterface.h"
+#include "llvm/include/llvm/ADT/STLExtras.h"  // from @llvm-project
+#include "mlir/include/mlir/IR/Operation.h"   // from @llvm-project
+#include "mlir/include/mlir/IR/Value.h"       // from @llvm-project
+#include "mlir/include/mlir/IR/Visitors.h"    // from @llvm-project
+#include "mlir/include/mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "mlir/include/mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace heir {
+
+bool isMutated(Value value) {
+  return llvm::any_of(value.getUsers(), [&](Operation* user) {
+    bool mutatesOperand = hasEffect<MemoryEffects::Write>(user, value);
+    auto inPlaceOp = dyn_cast<InPlaceOpInterface>(user);
+    bool isTargetOfInPlaceOp =
+        inPlaceOp &&
+        user->getOperand(inPlaceOp.getInPlaceOperandIndex()) == value;
+    return mutatesOperand || isTargetOfInPlaceOp;
+  });
+}
+
+ConstQualifierAnalysis::ConstQualifierAnalysis(Operation* op) {
+  op->walk([&](Operation* op) {
+    for (Value result : op->getResults()) {
+      if (!isMutated(result)) {
+        constValues.insert(result);
+      }
+    }
+  });
+}
+
+}  // namespace heir
+}  // namespace mlir

lib/Analysis/Cpp/ConstQualifierAnalysis.h

@@ -0,0 +1,30 @@
+#ifndef LIB_ANALYSIS_CPP_CONSTQUALIFIERANALYSIS_H_
+#define LIB_ANALYSIS_CPP_CONSTQUALIFIERANALYSIS_H_
+
+#include <cassert>
+
+#include "mlir/include/mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/include/mlir/IR/Value.h"      // from @llvm-project
+#include "mlir/include/mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace heir {
+
+class ConstQualifierAnalysis {
+ public:
+  ConstQualifierAnalysis(Operation* op);
+  ~ConstQualifierAnalysis() = default;
+
+  /// Return true if the value can be marked as const in C++ codegen.
+  bool canBeMarkedConst(Value value) const {
+    return constValues.contains(value);
+  }
+
+ private:
+  llvm::DenseSet<Value> constValues;
+};
+
+}  // namespace heir
+}  // namespace mlir
+
+#endif  // LIB_ANALYSIS_CPP_CONSTQUALIFIERANALYSIS_H_

lib/Dialect/Lattigo/IR/BUILD

@@ -29,7 +29,7 @@ cc_library(
         ":LattigoOps",
         ":LattigoTypes",
         "@heir//lib/Dialect:HEIRInterfaces",
-        "@heir//lib/Utils/Tablegen:InplaceOpInterface",
+        "@heir//lib/Utils/Tablegen:InPlaceOpInterface",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
     ],
@@ -88,7 +88,7 @@ cc_library(
         ":ops_inc_gen",
         ":types_inc_gen",
         "@heir//lib/Dialect:HEIRInterfaces",
-        "@heir//lib/Utils/Tablegen:InplaceOpInterface",
+        "@heir//lib/Utils/Tablegen:InPlaceOpInterface",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Support",
     ],

lib/Dialect/Lattigo/IR/LattigoBGVOps.td

@@ -43,7 +43,7 @@ def Lattigo_BGVNewEncoderOp : Lattigo_BGVOp<"new_encoder"> {
   let results = (outs Lattigo_BGVEncoder:$encoder);
 }
 
-def Lattigo_BGVEncodeOp : Lattigo_BGVOp<"encode", [InplaceOpInterface]> {
+def Lattigo_BGVEncodeOp : Lattigo_BGVOp<"encode", [InPlaceOpInterface]> {
   let summary = "Encode a plaintext value in the Lattigo BGV dialect";
   let description = [{
     This operation encodes a plaintext value using the specified encoder in the Lattigo BGV dialect.
@@ -60,7 +60,7 @@ def Lattigo_BGVEncodeOp : Lattigo_BGVOp<"encode", [InplaceOpInterface]> {
   );
   let results = (outs Lattigo_RLWEPlaintext:$encoded);
 
-  let extraClassDeclaration = "int getInplaceOperandIndex() { return 2; }";
+  let extraClassDeclaration = "int getInPlaceOperandIndex() { return 2; }";
 }
 
 def Lattigo_BGVDecodeOp : Lattigo_BGVOp<"decode", [AllTypesMatch<["value", "decoded"]>]> {
@@ -137,8 +137,8 @@ def Lattigo_BGVMulNewOp : Lattigo_BGVBinaryOp<"mul_new", [IncreasesMulDepthOpInt
   }];
 }
 
-class Lattigo_BGVBinaryInplaceOp<string mnemonic, list<Trait> traits = []> :
-        Lattigo_BGVOp<mnemonic, traits # [InplaceOpInterface]> {
+class Lattigo_BGVBinaryInPlaceOp<string mnemonic, list<Trait> traits = []> :
+        Lattigo_BGVOp<mnemonic, traits # [InPlaceOpInterface]> {
   let arguments = (ins
     Lattigo_BGVEvaluator:$evaluator,
     Lattigo_RLWECiphertext:$lhs,
@@ -149,10 +149,10 @@ class Lattigo_BGVBinaryInplaceOp<string mnemonic, list<Trait> traits = []> :
   );
   let results = (outs Lattigo_RLWECiphertext:$output);
 
-  let extraClassDeclaration = "int getInplaceOperandIndex() { return 3; }";
+  let extraClassDeclaration = "int getInPlaceOperandIndex() { return 3; }";
 }
 
-def Lattigo_BGVAddOp : Lattigo_BGVBinaryInplaceOp<"add"> {
+def Lattigo_BGVAddOp : Lattigo_BGVBinaryInPlaceOp<"add"> {
   let summary = "Add two ciphertexts in the Lattigo BGV dialect";
   let description = [{
     This operation adds two ciphertext values in the Lattigo BGV dialect.
@@ -162,7 +162,7 @@ def Lattigo_BGVAddOp : Lattigo_BGVBinaryInplaceOp<"add"> {
   }];
 }
 
-def Lattigo_BGVSubOp : Lattigo_BGVBinaryInplaceOp<"sub"> {
+def Lattigo_BGVSubOp : Lattigo_BGVBinaryInPlaceOp<"sub"> {
   let summary = "Subtract two ciphertexts in the Lattigo BGV dialect";
   let description = [{
     This operation subtracts one ciphertext value from another in the Lattigo BGV dialect.
@@ -172,7 +172,7 @@ def Lattigo_BGVSubOp : Lattigo_BGVBinaryInplaceOp<"sub"> {
   }];
 }
 
-def Lattigo_BGVMulOp : Lattigo_BGVBinaryInplaceOp<"mul", [IncreasesMulDepthOpInterface]> {
+def Lattigo_BGVMulOp : Lattigo_BGVBinaryInPlaceOp<"mul", [IncreasesMulDepthOpInterface]> {
   let summary = "Multiply two ciphertexts in the Lattigo BGV dialect";
   let description = [{
     This operation multiplies two ciphertext values in the Lattigo BGV dialect.
@@ -231,20 +231,20 @@ def Lattigo_BGVRotateRowsNewOp : Lattigo_BGVUnaryOp<"rotate_rows_new"> {
   }];
 }
 
-class Lattigo_BGVUnaryInplaceOp<string mnemonic> :
-        Lattigo_BGVOp<mnemonic, [InplaceOpInterface]> {
+class Lattigo_BGVUnaryInPlaceOp<string mnemonic> :
+        Lattigo_BGVOp<mnemonic, [InPlaceOpInterface]> {
   let arguments = (ins
     Lattigo_BGVEvaluator:$evaluator,
     Lattigo_RLWECiphertext:$input,
-    // see BinaryInplaceOp above
+    // see BinaryInPlaceOp above
     Lattigo_RLWECiphertext:$inplace
   );
   let results = (outs Lattigo_RLWECiphertext:$output);
 
-  let extraClassDeclaration = "int getInplaceOperandIndex() { return 2; }";
+  let extraClassDeclaration = "int getInPlaceOperandIndex() { return 2; }";
 }
 
-def Lattigo_BGVRelinearizeOp : Lattigo_BGVUnaryInplaceOp<"relinearize"> {
+def Lattigo_BGVRelinearizeOp : Lattigo_BGVUnaryInPlaceOp<"relinearize"> {
   let summary = "Relinearize a ciphertext in the Lattigo BGV dialect";
   let description = [{
     This operation relinearizes a ciphertext value in the Lattigo BGV dialect.
@@ -254,7 +254,7 @@ def Lattigo_BGVRelinearizeOp : Lattigo_BGVUnaryInplaceOp<"relinearize"> {
   }];
 }
 
-def Lattigo_BGVRescaleOp : Lattigo_BGVUnaryInplaceOp<"rescale"> {
+def Lattigo_BGVRescaleOp : Lattigo_BGVUnaryInPlaceOp<"rescale"> {
   let summary = "Rescale a ciphertext in the Lattigo BGV dialect";
   let description = [{
     This operation rescales a ciphertext value in the Lattigo BGV dialect.
@@ -264,7 +264,7 @@ def Lattigo_BGVRescaleOp : Lattigo_BGVUnaryInplaceOp<"rescale"> {
   }];
 }
 
-def Lattigo_BGVRotateColumnsOp : Lattigo_BGVUnaryInplaceOp<"rotate_columns"> {
+def Lattigo_BGVRotateColumnsOp : Lattigo_BGVUnaryInPlaceOp<"rotate_columns"> {
   let summary = "Rotate columns of a ciphertext in the Lattigo BGV dialect";
   let description = [{
     This operation rotates the columns of a ciphertext value in the Lattigo BGV dialect.
@@ -285,7 +285,7 @@ def Lattigo_BGVRotateColumnsOp : Lattigo_BGVUnaryInplaceOp<"rotate_columns"> {
   let results = (outs Lattigo_RLWECiphertext:$output);
 }
 
-def Lattigo_BGVRotateRowsOp : Lattigo_BGVUnaryInplaceOp<"rotate_rows"> {
+def Lattigo_BGVRotateRowsOp : Lattigo_BGVUnaryInPlaceOp<"rotate_rows"> {
   let summary = "Rotate rows of a ciphertext in the Lattigo BGV dialect";
   let description = [{
     This operation swap the rows of a ciphertext value in the Lattigo BGV dialect.

lib/Dialect/Lattigo/IR/LattigoCKKSOps.td

@@ -43,7 +43,7 @@ def Lattigo_CKKSNewEncoderOp : Lattigo_CKKSOp<"new_encoder"> {
   let results = (outs Lattigo_CKKSEncoder:$encoder);
 }
 
-def Lattigo_CKKSEncodeOp : Lattigo_CKKSOp<"encode", [InplaceOpInterface]> {
+def Lattigo_CKKSEncodeOp : Lattigo_CKKSOp<"encode", [InPlaceOpInterface]> {
   let summary = "Encode a plaintext value in the Lattigo CKKS dialect";
   let description = [{
     This operation encodes a plaintext value using the specified encoder in the Lattigo CKKS dialect.
@@ -60,7 +60,7 @@ def Lattigo_CKKSEncodeOp : Lattigo_CKKSOp<"encode", [InplaceOpInterface]> {
   );
   let results = (outs Lattigo_RLWEPlaintext:$encoded);
 
-  let extraClassDeclaration = "int getInplaceOperandIndex() { return 2; }";
+  let extraClassDeclaration = "int getInPlaceOperandIndex() { return 2; }";
 }
 
 def Lattigo_CKKSDecodeOp : Lattigo_CKKSOp<"decode", [AllTypesMatch<["value", "decoded"]>]> {
@@ -131,8 +131,8 @@ def Lattigo_CKKSMulNewOp : Lattigo_CKKSBinaryOp<"mul_new", [IncreasesMulDepthOpI
   }];
 }
 
-class Lattigo_CKKSBinaryInplaceOp<string mnemonic, list<Trait> traits = []> :
-        Lattigo_CKKSOp<mnemonic, traits # [InplaceOpInterface]> {
+class Lattigo_CKKSBinaryInPlaceOp<string mnemonic, list<Trait> traits = []> :
+        Lattigo_CKKSOp<mnemonic, traits # [InPlaceOpInterface]> {
   let arguments = (ins
     Lattigo_CKKSEvaluator:$evaluator,
     Lattigo_RLWECiphertext:$lhs,
@@ -143,10 +143,10 @@ class Lattigo_CKKSBinaryInplaceOp<string mnemonic, list<Trait> traits = []> :
   );
   let results = (outs Lattigo_RLWECiphertext:$output);
 
-  let extraClassDeclaration = "int getInplaceOperandIndex() { return 3; }";
+  let extraClassDeclaration = "int getInPlaceOperandIndex() { return 3; }";
 }
 
-def Lattigo_CKKSAddOp : Lattigo_CKKSBinaryInplaceOp<"add"> {
+def Lattigo_CKKSAddOp : Lattigo_CKKSBinaryInPlaceOp<"add"> {
   let summary = "Add two ciphertexts in the Lattigo CKKS dialect";
   let description = [{
     This operation adds two ciphertext values in the Lattigo CKKS dialect.
@@ -156,7 +156,7 @@ def Lattigo_CKKSAddOp : Lattigo_CKKSBinaryInplaceOp<"add"> {
   }];
 }
 
-def Lattigo_CKKSSubOp : Lattigo_CKKSBinaryInplaceOp<"sub"> {
+def Lattigo_CKKSSubOp : Lattigo_CKKSBinaryInPlaceOp<"sub"> {
   let summary = "Subtract two ciphertexts in the Lattigo CKKS dialect";
   let description = [{
     This operation subtracts one ciphertext value from another in the Lattigo CKKS dialect.
@@ -166,7 +166,7 @@ def Lattigo_CKKSSubOp : Lattigo_CKKSBinaryInplaceOp<"sub"> {
   }];
 }
 
-def Lattigo_CKKSMulOp : Lattigo_CKKSBinaryInplaceOp<"mul", [IncreasesMulDepthOpInterface]> {
+def Lattigo_CKKSMulOp : Lattigo_CKKSBinaryInPlaceOp<"mul", [IncreasesMulDepthOpInterface]> {
   let summary = "Multiply two ciphertexts in the Lattigo CKKS dialect";
   let description = [{
     This operation multiplies two ciphertext values in the Lattigo CKKS dialect.
@@ -218,20 +218,20 @@ def Lattigo_CKKSRotateNewOp : Lattigo_CKKSOp<"rotate_new"> {
   let results = (outs Lattigo_RLWECiphertext:$output);
 }
 
-class Lattigo_CKKSUnaryInplaceOp<string mnemonic> :
-        Lattigo_CKKSOp<mnemonic, [InplaceOpInterface]> {
+class Lattigo_CKKSUnaryInPlaceOp<string mnemonic> :
+        Lattigo_CKKSOp<mnemonic, [InPlaceOpInterface]> {
   let arguments = (ins
     Lattigo_CKKSEvaluator:$evaluator,
     Lattigo_RLWECiphertext:$input,
-    // see BinaryInplaceOp above
+    // see BinaryInPlaceOp above
     Lattigo_RLWECiphertext:$inplace
   );
   let results = (outs Lattigo_RLWECiphertext:$output);
 
-  let extraClassDeclaration = "int getInplaceOperandIndex() { return 2; }";
+  let extraClassDeclaration = "int getInPlaceOperandIndex() { return 2; }";
 }
 
-def Lattigo_CKKSRelinearizeOp : Lattigo_CKKSUnaryInplaceOp<"relinearize"> {
+def Lattigo_CKKSRelinearizeOp : Lattigo_CKKSUnaryInPlaceOp<"relinearize"> {
   let summary = "Relinearize a ciphertext in the Lattigo CKKS dialect";
   let description = [{
     This operation relinearizes a ciphertext value in the Lattigo CKKS dialect.
@@ -241,7 +241,7 @@ def Lattigo_CKKSRelinearizeOp : Lattigo_CKKSUnaryInplaceOp<"relinearize"> {
   }];
 }
 
-def Lattigo_CKKSRescaleOp : Lattigo_CKKSUnaryInplaceOp<"rescale"> {
+def Lattigo_CKKSRescaleOp : Lattigo_CKKSUnaryInPlaceOp<"rescale"> {
   let summary = "Rescale a ciphertext in the Lattigo CKKS dialect";
   let description = [{
     This operation rescales a ciphertext value in the Lattigo CKKS dialect.
@@ -251,7 +251,7 @@ def Lattigo_CKKSRescaleOp : Lattigo_CKKSUnaryInplaceOp<"rescale"> {
   }];
 }
 
-def Lattigo_CKKSRotateOp : Lattigo_CKKSUnaryInplaceOp<"rotate"> {
+def Lattigo_CKKSRotateOp : Lattigo_CKKSUnaryInPlaceOp<"rotate"> {
   let summary = "Rotate slots of a ciphertext in the Lattigo CKKS dialect";
   let description = [{
     This operation rotates slots of a ciphertext value in the Lattigo CKKS dialect.
@@ -268,7 +268,7 @@ def Lattigo_CKKSRotateOp : Lattigo_CKKSUnaryInplaceOp<"rotate"> {
   let arguments = (ins
     Lattigo_CKKSEvaluator:$evaluator,
     Lattigo_RLWECiphertext:$input,
-    // see BinaryInplaceOp above
+    // see BinaryInPlaceOp above
     Lattigo_RLWECiphertext:$inplace,
     Builtin_IntegerAttr:$offset
   );

lib/Dialect/Lattigo/IR/LattigoOps.h

@@ -5,7 +5,7 @@
 #include "lib/Dialect/HEIRInterfaces.h"
 #include "lib/Dialect/Lattigo/IR/LattigoDialect.h"
 #include "lib/Dialect/Lattigo/IR/LattigoTypes.h"
-#include "lib/Utils/Tablegen/InplaceOpInterface.h"
+#include "lib/Utils/Tablegen/InPlaceOpInterface.h"
 #include "mlir/include/mlir/IR/BuiltinOps.h"  // from @llvm-project
 // IWYU pragma: end_keep
 

lib/Dialect/Lattigo/IR/LattigoOps.td

@@ -4,7 +4,7 @@
 include "LattigoDialect.td"
 include "LattigoTypes.td"
 include "mlir/IR/OpBase.td"
-include "lib/Utils/Tablegen/InplaceOpInterface.td"
+include "lib/Utils/Tablegen/InPlaceOpInterface.td"
 
 class Lattigo_Op<string mnemonic, list<Trait> traits = []> :
         Op<Lattigo_Dialect, mnemonic, traits> {

lib/Dialect/Lattigo/IR/LattigoRLWEOps.td

@@ -132,7 +132,7 @@ def Lattigo_RLWEDropLevelNewOp : Lattigo_RLWEOp<"drop_level_new"> {
   let results = (outs Lattigo_RLWECiphertext:$output);
 }
 
-def Lattigo_RLWEDropLevelOp : Lattigo_RLWEOp<"drop_level", [InplaceOpInterface]> {
+def Lattigo_RLWEDropLevelOp : Lattigo_RLWEOp<"drop_level", [InPlaceOpInterface]> {
   let summary = "Drop level of a ciphertext";
   let description = [{
     This operation drops the level of a ciphertext
@@ -148,7 +148,7 @@ def Lattigo_RLWEDropLevelOp : Lattigo_RLWEOp<"drop_level", [InplaceOpInterface]>
   );
   let results = (outs Lattigo_RLWECiphertext:$output);
 
-  let extraClassDeclaration = "int getInplaceOperandIndex() { return 2; }";
+  let extraClassDeclaration = "int getInPlaceOperandIndex() { return 2; }";
 }
 
 def Lattigo_RLWENegateNewOp : Lattigo_RLWEOp<"negate_new"> {
@@ -160,7 +160,7 @@ def Lattigo_RLWENegateNewOp : Lattigo_RLWEOp<"negate_new"> {
   let results = (outs Lattigo_RLWECiphertext:$output);
 }
 
-def Lattigo_RLWENegateOp : Lattigo_RLWEOp<"negate", [InplaceOpInterface]> {
+def Lattigo_RLWENegateOp : Lattigo_RLWEOp<"negate", [InPlaceOpInterface]> {
   let summary = "Negate of a ciphertext";
   let description = [{
     This operation negates a ciphertext
@@ -175,7 +175,7 @@ def Lattigo_RLWENegateOp : Lattigo_RLWEOp<"negate", [InplaceOpInterface]> {
   );
   let results = (outs Lattigo_RLWECiphertext:$output);
 
-  let extraClassDeclaration = "int getInplaceOperandIndex() { return 2; }";
+  let extraClassDeclaration = "int getInPlaceOperandIndex() { return 2; }";
 }
 
 #endif  // LIB_DIALECT_LATTIGO_IR_LATTIGORLWEOPS_TD_