integrate split-preprocessing into pipelines

This includes a couple of minor changes:
* lwe-to-openfhe needs to update function signatures and call ops that have ciphertext data
* update split-preprocessing to batch plaintexts into groups with a configurable maximum # of return values

PiperOrigin-RevId: 872514667

Author: asraa

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

@@ -32,6 +32,7 @@
 #include "mlir/include/mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinTypes.h"           // from @llvm-project
 #include "mlir/include/mlir/IR/PatternMatch.h"           // from @llvm-project
+#include "mlir/include/mlir/IR/TypeUtilities.h"          // from @llvm-project
 #include "mlir/include/mlir/IR/Value.h"                  // from @llvm-project
 #include "mlir/include/mlir/Support/LLVM.h"              // from @llvm-project
 #include "mlir/include/mlir/Support/LogicalResult.h"     // from @llvm-project
@@ -114,10 +115,8 @@ bool containsArgumentOfDialect(Operation* op) {
     return false;
   }
   return llvm::any_of(funcOp.getArgumentTypes(), [&](Type argType) {
-    if (isa<ShapedType>(argType)) {
-      argType = cast<ShapedType>(argType).getElementType();
-    }
-    return DialectEqual<Dialects...>()(&argType.getDialect());
+    return DialectEqual<Dialects...>()(
+        &getElementTypeOrSelf(argType).getDialect());
   });
 }
 
@@ -794,7 +793,8 @@ struct LWEToLattigo : public impl::LWEToLattigoBase<LWEToLattigo> {
       return typeConverter.isSignatureLegal(op.getFunctionType()) &&
              typeConverter.isLegal(&op.getBody()) &&
              (!containsArgumentOfDialect<lwe::LWEDialect, bgv::BGVDialect,
-                                         ckks::CKKSDialect>(op) ||
+                                         ckks::CKKSDialect,
+                                         lattigo::LattigoDialect>(op) ||
               hasCryptoContextArg);
     });
 
@@ -809,7 +809,17 @@ struct LWEToLattigo : public impl::LWEToLattigoBase<LWEToLattigo> {
           !operandTypes.empty() &&
           mlir::isa<lattigo::BGVEvaluatorType, lattigo::CKKSEvaluatorType>(
               *operandTypes.begin());
-      return (!containsCryptoArg || hasCryptoContextArg);
+      // crypto context may need to be added for any function call whose callee
+      // has crypto ops.
+      auto containsCryptoOps = false;
+      FailureOr<func::FuncOp> callee = getCalledFunction(op);
+      if (succeeded(callee)) {
+        containsCryptoOps =
+            containsDialects<lwe::LWEDialect, bgv::BGVDialect,
+                             ckks::CKKSDialect, lattigo::LattigoDialect>(
+                callee.value());
+      }
+      return (!(containsCryptoArg || containsCryptoOps) || hasCryptoContextArg);
     });
 
     // All other operations are legal if they have no LWE typed operands or

lib/Dialect/LWE/Conversions/LWEToOpenfhe/LWEToOpenfhe.cpp

@@ -76,7 +76,8 @@ namespace {
 template <typename... Dialects>
 bool containsArgumentOfDialect(func::FuncOp funcOp) {
   return llvm::any_of(funcOp.getArgumentTypes(), [&](Type argType) {
-    return DialectEqual<Dialects...>()(&argType.getDialect());
+    return DialectEqual<Dialects...>()(
+        &getElementTypeOrSelf(argType).getDialect());
   });
 }
 
@@ -109,11 +110,12 @@ struct AddCryptoContextArg : public OpConversionPattern<func::FuncOp> {
     }
 
     auto containsCryptoOps =
-        containsDialects<lwe::LWEDialect, bgv::BGVDialect, ckks::CKKSDialect>(
-            op);
+        containsDialects<lwe::LWEDialect, bgv::BGVDialect, ckks::CKKSDialect,
+                         openfhe::OpenfheDialect>(op);
     auto containsCryptoArg =
         containsArgumentOfDialect<lwe::LWEDialect, bgv::BGVDialect,
-                                  ckks::CKKSDialect>(op);
+                                  ckks::CKKSDialect, openfhe::OpenfheDialect>(
+            op);
     if (!(containsCryptoOps || containsCryptoArg)) {
       return rewriter.notifyMatchFailure(
           op, "contains neither ops nor arg types from lwe/bgv/ckks dialects");
@@ -427,11 +429,12 @@ struct LWEToOpenfhe : public impl::LWEToOpenfheBase<LWEToOpenfhe> {
         return hasCryptoContextArg;
       }
       auto containsCryptoOps =
-          containsDialects<lwe::LWEDialect, bgv::BGVDialect, ckks::CKKSDialect>(
-              op);
+          containsDialects<lwe::LWEDialect, bgv::BGVDialect, ckks::CKKSDialect,
+                           openfhe::OpenfheDialect>(op);
       auto containsCryptoArg =
           containsArgumentOfDialect<lwe::LWEDialect, bgv::BGVDialect,
-                                    ckks::CKKSDialect>(op);
+                                    ckks::CKKSDialect, openfhe::OpenfheDialect>(
+              op);
       return typeConverter.isSignatureLegal(op.getFunctionType()) &&
              typeConverter.isLegal(&op.getBody()) &&
              (!(containsCryptoOps || containsCryptoArg) || hasCryptoContextArg);
@@ -447,11 +450,19 @@ struct LWEToOpenfhe : public impl::LWEToOpenfheBase<LWEToOpenfhe> {
         return hasCryptoContextArg;
       }
       auto containsCryptoArg = llvm::any_of(operandTypes, [&](Type argType) {
-        return DialectEqual<lwe::LWEDialect, bgv::BGVDialect,
-                            ckks::CKKSDialect>()(
+        return DialectEqual<lwe::LWEDialect, bgv::BGVDialect, ckks::CKKSDialect,
+                            openfhe::OpenfheDialect>()(
             &getElementTypeOrSelf(argType).getDialect());
       });
-      return (!containsCryptoArg || hasCryptoContextArg);
+      auto containsCryptoOps = false;
+      FailureOr<func::FuncOp> callee = getCalledFunction(op);
+      if (succeeded(callee)) {
+        containsCryptoOps =
+            containsDialects<lwe::LWEDialect, bgv::BGVDialect,
+                             ckks::CKKSDialect, openfhe::OpenfheDialect>(
+                callee.value());
+      }
+      return (!(containsCryptoArg || containsCryptoOps) || hasCryptoContextArg);
     });
 
     patterns.add<

lib/Pipelines/ArithmeticPipelineRegistration.cpp

@@ -186,8 +186,6 @@ void mlirToSecretArithmeticPipelineBuilder(
   // Balance Operations
   pm.addPass(createOperationBalancer());
 
-  // Add a __preprocessed helper for offline pre-packing of plaintexts
-  pm.addPass(createSplitPreprocessing());
   lowerAssignLayout(pm, false);
 
   // Add encrypt/decrypt helper functions for each function argument and return
@@ -406,6 +404,11 @@ void mlirToRLWEPipeline(OpPassManager& pm,
   pm.addPass(createCanonicalizerPass());
   pm.addPass(createCSEPass());
 
+  // Add a __preprocessed helper for offline pre-packing of plaintexts
+  auto splitPreprocessingOptions = SplitPreprocessingOptions{};
+  splitPreprocessingOptions.maxReturnValues = options.splitPreprocessing;
+  pm.addPass(createSplitPreprocessing(splitPreprocessingOptions));
+
   ElementwiseToAffineOptions elementwiseOptions;
   elementwiseOptions.convertDialects = {"ckks", "bgv", "lwe"};
   pm.addPass(createElementwiseToAffine(elementwiseOptions));
@@ -526,6 +529,7 @@ void torchLinalgToCkksBuilder(OpPassManager& manager,
   suboptions.ckksBootstrapWaterline = options.ckksBootstrapWaterline;
   suboptions.scalingModBits = options.scalingModBits;
   suboptions.firstModBits = options.firstModBits;
+  suboptions.splitPreprocessing = options.splitPreprocessing;
 
   mlirToRLWEPipelineBuilder(mlir::heir::RLWEScheme::ckksScheme)(manager,
                                                                 suboptions);

lib/Pipelines/ArithmeticPipelineRegistration.h

@@ -100,6 +100,11 @@ struct MlirToRLWEPipelineOptions : public SimdVectorizerOptions {
       llvm::cl::desc("File name to import execution result from (c.f. --secret-"
                      "import-execution-result)"),
       llvm::cl::init("")};
+  PassOptions::Option<int> splitPreprocessing{
+      *this, "split-preprocessing",
+      llvm::cl::desc("Split preprocessing into separate function with N return "
+                     "values (default to no split)"),
+      llvm::cl::init(0)};
 };
 
 struct PlaintextBackendOptions
@@ -179,6 +184,11 @@ struct TorchLinalgToCkksPipelineOptions
       llvm::cl::desc("The number of levels to keep until bootstrapping in CKKS "
                      "(c.f. --secret-insert-mgmt-ckks)"),
       llvm::cl::init(10)};
+  PassOptions::Option<int> splitPreprocessing{
+      *this, "split-preprocessing",
+      llvm::cl::desc("Split preprocessing into separate function with N return "
+                     "values (default to no split)"),
+      llvm::cl::init(0)};
 };
 void torchLinalgToCkksBuilder(OpPassManager& manager,
                               const TorchLinalgToCkksPipelineOptions& options);

lib/Transforms/FoldConstantTensors/FoldConstantTensors.cpp

@@ -210,7 +210,7 @@ class CollapseShapeAfterConstant final
       return rewriter.notifyMatchFailure(
           collapseOp, "source of collapse must be a constant");
 
-    auto sourceAttr = llvm::dyn_cast<ElementsAttr>(constantOp.getValue());
+    auto sourceAttr = llvm::dyn_cast<DenseElementsAttr>(constantOp.getValue());
     if (!sourceAttr)
       return rewriter.notifyMatchFailure(
           collapseOp, "source of collapse must be an elements attribute");