feat: Add secret-to-bgv conversion for ciphertext arith ops

Adds parameter setting options. All secret inputs are required to be a uniform tensor shape that matches the ring dimension in the parameters specified.

PiperOrigin-RevId: 615502361

include/Conversion/SecretToBGV/BUILD

@@ -0,0 +1,37 @@
+# SecretToBGV tablegen and headers.
+
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library")
+
+package(
+    default_applicable_licenses = ["@heir//:license"],
+    default_visibility = ["//visibility:public"],
+)
+
+exports_files(
+    [
+        "SecretToBGV.h",
+    ],
+)
+
+gentbl_cc_library(
+    name = "pass_inc_gen",
+    tbl_outs = [
+        (
+            [
+                "-gen-pass-decls",
+                "-name=SecretToBGV",
+            ],
+            "SecretToBGV.h.inc",
+        ),
+        (
+            ["-gen-pass-doc"],
+            "SecretToBGV.md",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "SecretToBGV.td",
+    deps = [
+        "@llvm-project//mlir:OpBaseTdFiles",
+        "@llvm-project//mlir:PassBaseTdFiles",
+    ],
+)

include/Conversion/SecretToBGV/SecretToBGV.h

@@ -0,0 +1,16 @@
+#ifndef INCLUDE_CONVERSION_SECRETTOBGV_SECRETTOBGV_H_
+#define INCLUDE_CONVERSION_SECRETTOBGV_SECRETTOBGV_H_
+
+#include "mlir/include/mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace mlir::heir {
+
+#define GEN_PASS_DECL
+#include "include/Conversion/SecretToBGV/SecretToBGV.h.inc"
+
+#define GEN_PASS_REGISTRATION
+#include "include/Conversion/SecretToBGV/SecretToBGV.h.inc"
+
+}  // namespace mlir::heir
+
+#endif  // INCLUDE_CONVERSION_SECRETTOBGV_SECRETTOBGV_H_

include/Conversion/SecretToBGV/SecretToBGV.td

@@ -0,0 +1,39 @@
+#ifndef INCLUDE_CONVERSION_SECRETTOBGV_SECRETTOBGV_TD_
+#define INCLUDE_CONVERSION_SECRETTOBGV_SECRETTOBGV_TD_
+
+include "mlir/Pass/PassBase.td"
+
+def SecretToBGV : Pass<"secret-to-bgv"> {
+  let summary = "Lower `secret` to `bgv` dialect.";
+
+  let description = [{
+    This pass lowers an IR with `secret.generic` blocks containing arithmetic
+    operations to operations on ciphertexts with the BGV dialect.
+
+    The pass assumes that the `secret.generic` regions have been distributed
+    through arithmetic operations so that only one ciphertext operation appears
+    per generic block. It also requires that `canonicalize` was run so that
+    non-secret values used are removed from the `secret.generic`'s block
+    arguments.
+
+    The pass requires that all types are tensors of a uniform shape matching the
+    dimension of the ciphertext space specified my `poly-mod-degree`.
+  }];
+
+  let dependentDialects = [
+    "mlir::heir::polynomial::PolynomialDialect",
+    "mlir::heir::bgv::BGVDialect",
+    "mlir::heir::lwe::LWEDialect",
+  ];
+
+  let options = [
+    Option<"polyModDegree", "poly-mod-degree", "int",
+           /*default=*/"1024", "Default degree of the cyclotomic polynomial "
+           "modulus to use for ciphertext space.">,
+    Option<"coefficientModBits", "coefficient-mod-bits", "int",
+           /*default=*/"29", "Default number of bits of the prime "
+           "coefficient modulus to use " "for the ciphertext space.">
+  ];
+}
+
+#endif  // INCLUDE_CONVERSION_SECRETTOBGV_SECRETTOBGV_TD_

include/Dialect/BGV/IR/BGVOps.td

@@ -27,7 +27,9 @@ class BGV_CiphertextPlaintextOp<string mnemonic, list<Trait> traits =
     [AllTypesMatch<["x", "output"]>,
     TypesMatchWith<"type of 'y' matches encoding type of 'x'",
                     "output", "y",
-                    "lwe::RLWEPlaintextType::get($_ctxt, ::llvm::cast<lwe::RLWECiphertextType>($_self).getEncoding())">]> :
+                    "lwe::RLWEPlaintextType::get($_ctxt, "
+                      "::llvm::cast<lwe::RLWECiphertextType>($_self).getEncoding(),"
+                      "::llvm::cast<lwe::RLWECiphertextType>($_self).getRlweParams().getRing())">]> :
         BGV_Op<mnemonic, traits> {
   let arguments = (ins
     RLWECiphertext:$x,

include/Dialect/LWE/IR/LWEAttributes.td

@@ -296,6 +296,8 @@ def RLWE_InverseCanonicalEmbeddingEncoding
   }];
 }
 
+def AnyRLWEEncodingAttr : AnyAttrOf<[RLWE_PolynomialCoefficientEncoding, RLWE_PolynomialEvaluationEncoding, RLWE_InverseCanonicalEmbeddingEncoding]>;
+
 def LWE_LWEParams : AttrDef<LWE_Dialect, "LWEParams"> {
   let mnemonic = "lwe_params";
 

include/Dialect/LWE/IR/LWEOps.td

@@ -3,6 +3,8 @@
 
 include "include/Dialect/LWE/IR/LWEDialect.td"
 include "include/Dialect/LWE/IR/LWETypes.td"
+include "include/Dialect/Polynomial/IR/PolynomialAttributes.td"
+
 include "mlir/IR/BuiltinAttributeInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
@@ -24,7 +26,7 @@ def LWE_EncodeOp : LWE_Op<"encode", [Pure]> {
     Examples:
 
     ```
-    %Y = lwe.encode %value {encoding = #enc}: i1 to !lwe.lwe_plaintext<encoding = #enc, ring = #ring>
+    %Y = lwe.encode %value {encoding = #enc}: i1 to !lwe.lwe_plaintext<encoding = #enc>
     ```
   }];
 
@@ -59,4 +61,33 @@ def LWE_TrivialEncryptOp: LWE_Op<"trivial_encrypt", [Pure]> {
   let hasVerifier = 1;
 }
 
+def LWE_RLWEEncodeOp : LWE_Op<"rlwe_encode", [Pure]> {
+  let summary = "Encode an integer to yield an RLWE plaintext";
+  let description = [{
+    Encode an integer to yield an RLWE plaintext.
+
+    This op uses a an encoding attribute to encode the bits of the integer into
+    an RLWE plaintext value that can then be encrypted.
+
+    Examples:
+
+    ```
+    %Y = lwe.rlwe_encode %value {encoding = #enc, ring = #ring}: i1 to !lwe.rlwe_plaintext<encoding = #enc, ring = #ring>
+    ```
+  }];
+
+  let arguments = (ins
+    SignlessIntegerLike:$plaintext,
+    AnyRLWEEncodingAttr:$encoding,
+    Ring_Attr:$ring
+  );
+
+  let results = (outs RLWEPlaintext:$output);
+  let assemblyFormat = "$plaintext attr-dict `:` qualified(type($plaintext)) `to` qualified(type($output))";
+
+  // Verify that the encoding and ring parameter matches the output plaintext attribute.
+  let hasVerifier = 1;
+}
+
+
 #endif  // HEIR_INCLUDE_DIALECT_LWE_IR_LWEOPS_TD_

include/Dialect/LWE/IR/LWETypes.td

@@ -43,7 +43,7 @@ def RLWECiphertext : LWE_Type<"RLWECiphertext", "rlwe_ciphertext"> {
 
   let parameters = (ins
     "::mlir::Attribute":$encoding,
-    OptionalParameter<"RLWEParamsAttr">:$rlwe_params
+    "RLWEParamsAttr":$rlwe_params
   );
 
   let assemblyFormat = "`<` struct(params) `>`";
@@ -70,7 +70,8 @@ def RLWEPlaintext : LWE_Type<"RLWEPlaintext", "rlwe_plaintext"> {
   let summary = "A type for RLWE plaintexts";
 
   let parameters = (ins
-    "::mlir::Attribute":$encoding
+    "::mlir::Attribute":$encoding,
+    "::mlir::heir::polynomial::RingAttr":$ring
   );
 
   let assemblyFormat = "`<` struct(params) `>`";

lib/Conversion/SecretToBGV/BUILD

@@ -0,0 +1,28 @@
+package(
+    default_applicable_licenses = ["@heir//:license"],
+    default_visibility = ["//visibility:public"],
+)
+
+cc_library(
+    name = "SecretToBGV",
+    srcs = ["SecretToBGV.cpp"],
+    hdrs = [
+        "@heir//include/Conversion/SecretToBGV:SecretToBGV.h",
+    ],
+    deps = [
+        "@heir//include/Conversion/SecretToBGV:pass_inc_gen",
+        "@heir//lib/Conversion:Utils",
+        "@heir//lib/Dialect/BGV/IR:Dialect",
+        "@heir//lib/Dialect/LWE/IR:Dialect",
+        "@heir//lib/Dialect/Polynomial/IR:Dialect",
+        "@heir//lib/Dialect/Polynomial/IR:Polynomial",
+        "@heir//lib/Dialect/Polynomial/IR:PolynomialAttributes",
+        "@heir//lib/Dialect/Secret/IR:Dialect",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
+    ],
+)