Merge pull request google#892 from AlexanderViand-Intel:client-interf…

…ace-rename

PiperOrigin-RevId: 662193084

lib/Dialect/BGV/CMakeLists.txt

@@ -1,16 +1,13 @@
 add_subdirectory(IR)
-add_subdirectory(Transforms)
 add_mlir_dialect_library(MLIRBGV
         IR/BGVDialect.cpp
-        Transforms/AddClientInterface.cpp
 
         ADDITIONAL_HEADER_DIRS
         ${CMAKE_CURRENT_SOURCE_DIR}/IR
 
         DEPENDS
         MLIRBGVIncGen
         MLIRBGVOpsIncGen
-        MLIRBGVPassesIncGen
 
         LINK_LIBS PUBLIC
         MLIRIR

lib/Dialect/LWE/CMakeLists.txt

@@ -3,6 +3,7 @@ add_subdirectory(Transforms)
 add_mlir_dialect_library(MLIRLWE
         IR/LWEDialect.cpp
         Transforms/SetDefaultParameters.cpp
+        Transforms/AddClientInterface.cpp
 
         ADDITIONAL_HEADER_DIRS
         ${CMAKE_CURRENT_SOURCE_DIR}/IR

lib/Dialect/BGV/Transforms/AddClientInterface.cpp → lib/Dialect/LWE/Transforms/AddClientInterface.cpp

@@ -1,9 +1,8 @@
-#include "lib/Dialect/BGV/Transforms/AddClientInterface.h"
+#include "lib/Dialect/LWE/Transforms/AddClientInterface.h"
 
 #include <cstddef>
 #include <string>
 
-#include "lib/Dialect/BGV/IR/BGVOps.h"
 #include "lib/Dialect/LWE/IR/LWEAttributes.h"
 #include "lib/Dialect/LWE/IR/LWEOps.h"
 #include "lib/Dialect/LWE/IR/LWETypes.h"
@@ -21,10 +20,10 @@
 
 namespace mlir {
 namespace heir {
-namespace bgv {
+namespace lwe {
 
 #define GEN_PASS_DEF_ADDCLIENTINTERFACE
-#include "lib/Dialect/BGV/Transforms/Passes.h.inc"
+#include "lib/Dialect/LWE/Transforms/Passes.h.inc"
 
 FailureOr<lwe::RLWEParamsAttr> getRlweParmsFromFuncOp(func::FuncOp op) {
   lwe::RLWEParamsAttr rlweParams = nullptr;
@@ -154,6 +153,7 @@ LogicalResult convertFunc(func::FuncOp op, bool usePublicKey,
   auto module = op->getParentOfType<ModuleOp>();
   auto rlweParamsResult = getRlweParmsFromFuncOp(op);
   if (failed(rlweParamsResult)) {
+    // TODO (#891): Add support for schemes other than BGV
     return failure();
   }
   lwe::RLWEParamsAttr rlweParams = rlweParamsResult.value();
@@ -271,6 +271,6 @@ struct AddClientInterface : impl::AddClientInterfaceBase<AddClientInterface> {
     if (result.wasInterrupted()) signalPassFailure();
   }
 };
-}  // namespace bgv
+}  // namespace lwe
 }  // namespace heir
 }  // namespace mlir

lib/Dialect/LWE/Transforms/AddClientInterface.h

@@ -0,0 +1,17 @@
+#ifndef LIB_DIALECT_LWE_TRANSFORMS_ADDCLIENTINTERFACE_H_
+#define LIB_DIALECT_LWE_TRANSFORMS_ADDCLIENTINTERFACE_H_
+
+#include "mlir/include/mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace mlir {
+namespace heir {
+namespace lwe {
+
+#define GEN_PASS_DECL_ADDCLIENTINTERFACE
+#include "lib/Dialect/LWE/Transforms/Passes.h.inc"
+
+}  // namespace lwe
+}  // namespace heir
+}  // namespace mlir
+
+#endif  // LIB_DIALECT_LWE_TRANSFORMS_ADDCLIENTINTERFACE_H_

lib/Dialect/LWE/Transforms/BUILD

@@ -11,12 +11,31 @@ cc_library(
         "Passes.h",
     ],
     deps = [
+        ":AddClientInterface",
         ":SetDefaultParameters",
         ":pass_inc_gen",
         "@heir//lib/Dialect/LWE/IR:Dialect",
     ],
 )
 
+cc_library(
+    name = "AddClientInterface",
+    srcs = ["AddClientInterface.cpp"],
+    hdrs = [
+        "AddClientInterface.h",
+    ],
+    deps = [
+        ":pass_inc_gen",
+        "@heir//lib/Dialect/BGV/IR:Dialect",
+        "@heir//lib/Dialect/LWE/IR:Dialect",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
 cc_library(
     name = "SetDefaultParameters",
     srcs = ["SetDefaultParameters.cpp"],

lib/Dialect/LWE/Transforms/Passes.h

@@ -2,6 +2,7 @@
 #define LIB_DIALECT_LWE_TRANSFORMS_PASSES_H_
 
 #include "lib/Dialect/LWE/IR/LWEDialect.h"
+#include "lib/Dialect/LWE/Transforms/AddClientInterface.h"
 #include "lib/Dialect/LWE/Transforms/SetDefaultParameters.h"
 
 namespace mlir {

lib/Dialect/LWE/Transforms/Passes.td

@@ -3,6 +3,25 @@
 
 include "mlir/Pass/PassBase.td"
 
+
+def AddClientInterface : Pass<"lwe-add-client-interface"> {
+  let summary = "Add client interfaces to (R)LWE encrypted functions";
+  let description = [{
+  This pass adds encrypt and decrypt functions for each compiled function in the
+  IR. These functions maintain the same interface as the original function,
+  while the compiled function may lose some of this information by the lowerings
+  to ciphertext types (e.g., a scalar ciphertext, when lowered through RLWE schemes,
+  must be encoded as a tensor).
+  }];
+  let dependentDialects = ["mlir::heir::lwe::LWEDialect"];
+  let options = [
+    Option<"usePublicKey", "use-public-key", "bool", /*default=*/"false",
+           "If true, generate a client interface that uses a public key for encryption.">,
+    Option<"oneValuePerHelperFn", "one-value-per-helper-fn", "bool", /*default=*/"false",
+           "If true, split encryption helpers into separate functions for each SSA value.">
+  ];
+}
+
 def SetDefaultParameters : Pass<"lwe-set-default-parameters"> {
   let summary = "Set default parameters for LWE ops";
   let description = [{

tests/bgv/add_client_interface.mlir → tests/lwe/add_client_interface.mlir

@@ -1,4 +1,4 @@
-// RUN: heir-opt --bgv-add-client-interface %s | FileCheck %s
+// RUN: heir-opt --lwe-add-client-interface %s | FileCheck %s
 
 // These two types differ only on their underlying_type. The IR stays as the !in_ty
 // for the entire computation until the final extract op.

tests/bgv/add_client_interface_public_key.mlir → tests/lwe/add_client_interface_public_key.mlir

@@ -1,4 +1,4 @@
-// RUN: heir-opt --bgv-add-client-interface=use-public-key=true %s | FileCheck %s
+// RUN: heir-opt --lwe-add-client-interface=use-public-key=true %s | FileCheck %s
 
 // These two types differ only on their underlying_type. The IR stays as the !in_ty
 // for the entire computation until the final extract op.

tests/bgv/add_client_interface_split.mlir → tests/lwe/add_client_interface_split.mlir

@@ -1,4 +1,4 @@
-// RUN: heir-opt '--bgv-add-client-interface=use-public-key=true one-value-per-helper-fn=true' %s | FileCheck %s
+// RUN: heir-opt '--lwe-add-client-interface=use-public-key=true one-value-per-helper-fn=true' %s | FileCheck %s
 
 #encoding = #lwe.polynomial_evaluation_encoding<cleartext_start = 16, cleartext_bitwidth = 16>
 #params = #lwe.rlwe_params<ring = <coefficientType = i32, coefficientModulus = 463187969 : i32, polynomialModulus=#polynomial.int_polynomial<1 + x**8>>>

tools/BUILD

@@ -46,15 +46,14 @@ cc_binary(
         "@heir//lib/Conversion/SecretToBGV",
         "@heir//lib/Dialect/ArithExt/IR:Dialect",
         "@heir//lib/Dialect/BGV/IR:Dialect",
-        "@heir//lib/Dialect/BGV/Transforms",
-        "@heir//lib/Dialect/BGV/Transforms:AddClientInterface",
         "@heir//lib/Dialect/CGGI/IR:Dialect",
         "@heir//lib/Dialect/CGGI/Transforms",
         "@heir//lib/Dialect/CKKS/IR:Dialect",
         "@heir//lib/Dialect/Comb/IR:Dialect",
         "@heir//lib/Dialect/Jaxite/IR:Dialect",
         "@heir//lib/Dialect/LWE/IR:Dialect",
         "@heir//lib/Dialect/LWE/Transforms",
+        "@heir//lib/Dialect/LWE/Transforms:AddClientInterface",
         "@heir//lib/Dialect/Openfhe/IR:Dialect",
         "@heir//lib/Dialect/Openfhe/Transforms",
         "@heir//lib/Dialect/Openfhe/Transforms:ConfigureCryptoContext",