openfhe emitter: add support for CKKS scheme emission

* Adds an end to end 16x16 matrix multiplication - Halevi-Shoup kernels haven't been implemented yet, so this is a naive implementation (one element of a vector per ciphertext, so not even packing the input at all) - it takes about 3 sec to run with OpenFHE * Adds option --openfhe-scheme={bgv,ckks} to heir-translate to indicate which API to use * Adds tensor operations in heir-translate PiperOrigin-RevId: 683254954
j2kun · Oct 7, 2024 · ad8e981 · ad8e981
1 parent 9d3391f
commit ad8e981
Show file tree

Hide file tree

Showing 24 changed files with 626 additions and 154 deletions.
diff --git a/lib/Analysis/SelectVariableNames/BUILD b/lib/Analysis/SelectVariableNames/BUILD
@@ -10,6 +10,9 @@ cc_library(
     hdrs = ["SelectVariableNames.h"],
     deps = [
         "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
     ],
 )
diff --git a/lib/Analysis/SelectVariableNames/SelectVariableNames.cpp b/lib/Analysis/SelectVariableNames/SelectVariableNames.cpp
@@ -4,6 +4,7 @@
 #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/Support/LLVM.h"  // from @llvm-project
 
 namespace mlir {
 namespace heir {

diff --git a/lib/Analysis/SelectVariableNames/SelectVariableNames.h b/lib/Analysis/SelectVariableNames/SelectVariableNames.h
@@ -1,11 +1,14 @@
 #ifndef LIB_ANALYSIS_SELECTVARIABLENAMES_SELECTVARIABLENAMES_H_
 #define LIB_ANALYSIS_SELECTVARIABLENAMES_SELECTVARIABLENAMES_H_
 
+#include <cassert>
 #include <string>
 
-#include "llvm/include/llvm/ADT/DenseMap.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 "llvm/include/llvm/ADT/DenseMap.h"          // from @llvm-project
+#include "mlir/include/mlir/IR/BuiltinAttributes.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/Support/LLVM.h"          // from @llvm-project
 
 namespace mlir {
 namespace heir {

diff --git a/lib/Conversion/SecretToCKKS/SecretToCKKS.cpp b/lib/Conversion/SecretToCKKS/SecretToCKKS.cpp
@@ -225,6 +225,7 @@ struct SecretToCKKS : public impl::SecretToCKKSBase<SecretToCKKS> {
         SecretGenericOpCipherConversion<arith::SubIOp, ckks::SubOp>,
         SecretGenericOpCipherConversion<arith::AddFOp, ckks::AddOp>,
         SecretGenericOpCipherConversion<arith::SubFOp, ckks::SubOp>,
+        SecretGenericOpCipherConversion<tensor::EmptyOp, tensor::EmptyOp>,
         SecretGenericTensorExtractConversion,
         SecretGenericTensorInsertConversion,
         SecretGenericOpRotateConversion<ckks::RotateOp>,

diff --git a/lib/Target/OpenFhePke/BUILD b/lib/Target/OpenFhePke/BUILD
@@ -5,10 +5,38 @@ package(
     default_visibility = ["//visibility:public"],
 )
 
+cc_library(
+    name = "OpenFheRegistration",
+    srcs = [
+        "OpenFheTranslateRegistration.cpp",
+    ],
+    hdrs = [
+        "OpenFheTranslateRegistration.h",
+    ],
+    deps = [
+        ":OpenFhePkeEmitter",
+        ":OpenFhePkeHeaderEmitter",
+        ":OpenFheUtils",
+        "@heir//lib/Analysis/SelectVariableNames",
+        "@heir//lib/Dialect/LWE/IR:Dialect",
+        "@heir//lib/Dialect/Openfhe/IR:Dialect",
+        "@heir//lib/Target:Utils",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:PolynomialDialect",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:TranslateLib",
+    ],
+)
+
 cc_library(
     name = "OpenFheUtils",
     srcs = ["OpenFheUtils.cpp"],
     hdrs = [
+        "OpenFhePkeTemplates.h",
         "OpenFheUtils.h",
     ],
     deps = [
@@ -26,7 +54,6 @@ cc_library(
     srcs = ["OpenFhePkeEmitter.cpp"],
     hdrs = [
         "OpenFhePkeEmitter.h",
-        "OpenFhePkeTemplates.h",
     ],
     deps = [
         ":OpenFheUtils",
@@ -38,10 +65,8 @@ cc_library(
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
-        "@llvm-project//mlir:PolynomialDialect",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TensorDialect",
-        "@llvm-project//mlir:TranslateLib",
     ],
 )
 
@@ -55,15 +80,12 @@ cc_library(
     deps = [
         ":OpenFheUtils",
         "@heir//lib/Analysis/SelectVariableNames",
-        "@heir//lib/Dialect/LWE/IR:Dialect",
         "@heir//lib/Dialect/Openfhe/IR:Dialect",
         "@heir//lib/Target:Utils",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
-        "@llvm-project//mlir:PolynomialDialect",
         "@llvm-project//mlir:Support",
-        "@llvm-project//mlir:TranslateLib",
     ],
 )