split-preprocessing: use new preprocessing dialect infrastructure

This change modifies split-preprocessing to use the new preprocessing dialect for its output, and updates the pipelines to use preprocessing-to-memref.

It changes the `split-preprocessing` flag from a numeric batching factor to a boolean enable/disable. It also adds a small amount of fixes to the lattigo lowering (which otherwise supports bufferization and memrefs).

Fixes #2960

PiperOrigin-RevId: 938040235

Author: j2kun

lib/Dialect/Preprocessing/Conversions/PreprocessingToLattigo/BUILD

@@ -16,6 +16,7 @@ cc_library(
         "@heir//lib/Dialect/Lattigo/IR:Dialect",
         "@heir//lib/Dialect/Preprocessing/Conversions:Util",
         "@heir//lib/Dialect/Preprocessing/IR:Dialect",
+        "@heir//lib/Utils",
         "@heir//lib/Utils:ConversionUtils",
         "@llvm-project//mlir:AffineDialect",
         "@llvm-project//mlir:ArithDialect",

lib/Dialect/Preprocessing/Conversions/PreprocessingToLattigo/PreprocessingToLattigo.cpp

@@ -8,6 +8,7 @@
 #include "lib/Dialect/Preprocessing/Conversions/Util.h"
 #include "lib/Dialect/Preprocessing/IR/PreprocessingDialect.h"
 #include "lib/Utils/ConversionUtils.h"
+#include "lib/Utils/Utils.h"
 #include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"   // from @llvm-project
 #include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
@@ -31,6 +32,10 @@ struct PreprocessingToLattigo
   void runOnOperation() override {
     ModuleOp module = getOperation();
 
+    if (!containsDialects<PreprocessingDialect>(module)) {
+      return;
+    }
+
     PreprocessingStorageLayoutAnalysis analysis(module);
     if (!analysis.isValid()) {
       signalPassFailure();

lib/Dialect/Preprocessing/Conversions/PreprocessingToMemref/BUILD

@@ -15,6 +15,7 @@ cc_library(
         "@heir//lib/Analysis/PreprocessingStorageLayoutAnalysis",
         "@heir//lib/Dialect/Preprocessing/Conversions:Util",
         "@heir//lib/Dialect/Preprocessing/IR:Dialect",
+        "@heir//lib/Utils",
         "@heir//lib/Utils:ConversionUtils",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:AffineAnalysis",

lib/Dialect/Preprocessing/Conversions/PreprocessingToMemref/PreprocessingToMemref.cpp

@@ -13,6 +13,7 @@
 #include "lib/Dialect/Preprocessing/IR/PreprocessingOps.h"
 #include "lib/Dialect/Preprocessing/IR/PreprocessingTypes.h"
 #include "lib/Utils/ConversionUtils.h"
+#include "lib/Utils/Utils.h"
 #include "llvm/include/llvm/ADT/STLExtras.h"    // from @llvm-project
 #include "llvm/include/llvm/ADT/SmallVector.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h"  // from @llvm-project
@@ -214,6 +215,10 @@ struct PreprocessingToMemref
   void runOnOperation() override {
     ModuleOp module = getOperation();
 
+    if (!containsDialects<PreprocessingDialect>(module)) {
+      return;
+    }
+
     PreprocessingStorageLayoutAnalysis analysis(module);
     if (!analysis.isValid()) {
       signalPassFailure();

lib/Dialect/Preprocessing/Conversions/PreprocessingToOpenfhe/BUILD

@@ -16,6 +16,7 @@ cc_library(
         "@heir//lib/Dialect/Openfhe/IR:Dialect",
         "@heir//lib/Dialect/Preprocessing/Conversions:Util",
         "@heir//lib/Dialect/Preprocessing/IR:Dialect",
+        "@heir//lib/Utils",
         "@heir//lib/Utils:ConversionUtils",
         "@llvm-project//mlir:AffineDialect",
         "@llvm-project//mlir:ArithDialect",

lib/Dialect/Preprocessing/Conversions/PreprocessingToOpenfhe/PreprocessingToOpenfhe.cpp

@@ -8,6 +8,7 @@
 #include "lib/Dialect/Preprocessing/Conversions/Util.h"
 #include "lib/Dialect/Preprocessing/IR/PreprocessingDialect.h"
 #include "lib/Utils/ConversionUtils.h"
+#include "lib/Utils/Utils.h"
 #include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"   // from @llvm-project
 #include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
@@ -31,6 +32,10 @@ struct PreprocessingToOpenfhe
   void runOnOperation() override {
     ModuleOp module = getOperation();
 
+    if (!containsDialects<PreprocessingDialect>(module)) {
+      return;
+    }
+
     PreprocessingStorageLayoutAnalysis analysis(module);
     if (!analysis.isValid()) {
       signalPassFailure();

lib/Dialect/Preprocessing/IR/PreprocessingOps.td

@@ -14,7 +14,7 @@ class Preprocessing_Op<string mnemonic, list<Trait> traits = []> :
   let cppNamespace = "::mlir::heir::preprocessing";
 }
 
-def Preprocessing_EmptyOp : Preprocessing_Op<"empty"> {
+def Preprocessing_EmptyOp : Preprocessing_Op<"empty", [MemoryEffectsOpInterface]> {
   let summary = "Allocate an empty preprocessing storage";
   let description = [{
     This op creates a new `preprocessing.storage` value with an undetermined size.
@@ -23,7 +23,7 @@ def Preprocessing_EmptyOp : Preprocessing_Op<"empty"> {
   let assemblyFormat = "attr-dict `:` type($storage)";
 }
 
-def Preprocessing_StoreOp : Preprocessing_Op<"store"> {
+def Preprocessing_StoreOp : Preprocessing_Op<"store", [MemoryEffectsOpInterface]> {
   let summary = "Store a value to preprocessing storage at given indices and site_id";
   let description = [{
     This op stores a value into a `preprocessing.storage`, while tracking:
@@ -49,7 +49,7 @@ def Preprocessing_StoreOp : Preprocessing_Op<"store"> {
   let assemblyFormat = "$value `,` $storage `[` $indices `]` `site` $site_id `<` $element_type `>` attr-dict `:` type($value) `,` type($storage)";
 }
 
-def Preprocessing_LoadOp : Preprocessing_Op<"load"> {
+def Preprocessing_LoadOp : Preprocessing_Op<"load", [MemoryEffectsOpInterface]> {
   let summary = "Load a value from preprocessing storage at given indices and site_id";
   let description = [{
     This op loads a value from a `preprocessing.storage`, while tracking:

lib/Dialect/Preprocessing/Transforms/BUILD

@@ -23,6 +23,7 @@ cc_library(
         ":pass_inc_gen",
         "@heir//lib/Dialect/Preprocessing/IR:Dialect",
         "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:Support",

lib/Dialect/Preprocessing/Transforms/ValidatePreprocessing.cpp

@@ -6,12 +6,13 @@
 
 #include "lib/Dialect/Preprocessing/IR/PreprocessingOps.h"
 #include "lib/Dialect/Preprocessing/IR/PreprocessingTypes.h"
-#include "llvm/include/llvm/ADT/ArrayRef.h"     // from @llvm-project
-#include "llvm/include/llvm/ADT/DenseMap.h"     // from @llvm-project
-#include "llvm/include/llvm/ADT/SmallVector.h"  // from @llvm-project
-#include "mlir/include/mlir/IR/Diagnostics.h"   // from @llvm-project
-#include "mlir/include/mlir/IR/Visitors.h"      // from @llvm-project
-#include "mlir/include/mlir/Support/LLVM.h"     // from @llvm-project
+#include "llvm/include/llvm/ADT/ArrayRef.h"             // from @llvm-project
+#include "llvm/include/llvm/ADT/DenseMap.h"             // from @llvm-project
+#include "llvm/include/llvm/ADT/SmallVector.h"          // from @llvm-project
+#include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/include/mlir/IR/Diagnostics.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 {
@@ -25,30 +26,36 @@ struct ValidatePreprocessing
   using ValidatePreprocessingBase::ValidatePreprocessingBase;
 
   void runOnOperation() override {
-    SmallVector<EmptyOp, 2> emptyOps;
+    Operation* module = getOperation();
+
+    bool hasMultipleEmpties = false;
+    module->walk([&](func::FuncOp funcOp) {
+      SmallVector<EmptyOp> emptyOps(funcOp.getOps<EmptyOp>());
+      if (emptyOps.size() > 1) {
+        for (size_t i = 1; i < emptyOps.size(); ++i) {
+          auto diag =
+              emptyOps[i]->emitOpError()
+              << "more than one preprocessing.empty allocation in function";
+          diag.attachNote(emptyOps[0]->getLoc()) << "previous allocation here";
+        }
+        hasMultipleEmpties = true;
+      }
+    });
+    if (hasMultipleEmpties) {
+      signalPassFailure();
+    }
+
     DenseMap<uint32_t, SmallVector<StoreOp, 2>> storesBySite;
     DenseMap<uint32_t, SmallVector<LoadOp, 2>> loadsBySite;
 
-    getOperation()->walk([&](Operation* op) {
-      if (auto emptyOp = dyn_cast<EmptyOp>(op)) {
-        emptyOps.push_back(emptyOp);
-      } else if (auto storeOp = dyn_cast<StoreOp>(op)) {
+    module->walk([&](Operation* op) {
+      if (auto storeOp = dyn_cast<StoreOp>(op)) {
         storesBySite[storeOp.getSiteId()].push_back(storeOp);
       } else if (auto loadOp = dyn_cast<LoadOp>(op)) {
         loadsBySite[loadOp.getSiteId()].push_back(loadOp);
       }
     });
 
-    if (emptyOps.size() > 1) {
-      for (size_t i = 1; i < emptyOps.size(); ++i) {
-        auto diag = emptyOps[i]->emitOpError()
-                    << "more than one preprocessing.empty allocation in module";
-        diag.attachNote(emptyOps[0]->getLoc()) << "previous allocation here";
-      }
-      signalPassFailure();
-      // Do NOT return, continue to site pairing checks
-    }
-
     std::set<uint32_t> siteIds;
     for (const auto& [siteId, stores] : storesBySite) {
       siteIds.insert(siteId);

lib/Pipelines/ArithmeticPipelineRegistration.cpp

@@ -16,6 +16,9 @@
 #include "lib/Dialect/Openfhe/Transforms/ConfigureCryptoContext.h"
 #include "lib/Dialect/Openfhe/Transforms/CountAddAndKeySwitch.h"
 #include "lib/Dialect/Openfhe/Transforms/FastRotationPrecompute.h"
+#include "lib/Dialect/Preprocessing/Conversions/PreprocessingToLattigo/PreprocessingToLattigo.h"
+#include "lib/Dialect/Preprocessing/Conversions/PreprocessingToOpenfhe/PreprocessingToOpenfhe.h"
+#include "lib/Dialect/Preprocessing/Transforms/ValidatePreprocessing.h"
 #include "lib/Dialect/Secret/Conversions/SecretToBGV/SecretToBGV.h"
 #include "lib/Dialect/Secret/Conversions/SecretToCKKS/SecretToCKKS.h"
 #include "lib/Dialect/Secret/Conversions/SecretToModArith/SecretToModArith.h"
@@ -142,6 +145,8 @@ void cleanupAfterLowerAssignLayout(OpPassManager& pm) {
   pm.addPass(createCSEPass());
   pm.addPass(createRemoveDeadValuesPass());
   pm.addPass(createSymbolDCEPass());
+  pm.addPass(createCanonicalizerPass());
+  pm.addPass(createCSEPass());
 }
 
 // Implement layout conversions as shift networks
@@ -423,9 +428,10 @@ void mlirToRLWEPipeline(OpPassManager& pm,
   pm.addPass(lwe::createImplementTrivialEncryptionAsAddition());
 
   // Add a __preprocessed helper for offline pre-packing of plaintexts
-  auto splitPreprocessingOptions = SplitPreprocessingOptions{};
-  splitPreprocessingOptions.maxReturnValues = options.splitPreprocessing;
-  pm.addPass(createSplitPreprocessing(splitPreprocessingOptions));
+  if (options.enableSplitPreprocessing) {
+    pm.addPass(createSplitPreprocessing());
+    pm.addPass(preprocessing::createValidatePreprocessing());
+  }
 
   ElementwiseToAffineOptions elementwiseOptions;
   elementwiseOptions.convertDialects = {"ckks", "bgv", "lwe"};
@@ -471,6 +477,7 @@ BackendPipelineBuilder toOpenFhePipelineBuilder() {
 
     // Convert LWE (and scheme-specific CKKS/BGV ops) to OpenFHE
     pm.addPass(lwe::createLWEToOpenfhe());
+    pm.addPass(preprocessing::createPreprocessingToOpenfhe());
     pm.addPass(createCanonicalizerPass());
     pm.addPass(createCSEPass());
 
@@ -513,6 +520,7 @@ BackendPipelineBuilder toLattigoPipelineBuilder() {
 
     // Convert LWE (and scheme-specific BGV ops) to Lattigo
     pm.addPass(lwe::createLWEToLattigo());
+    pm.addPass(preprocessing::createPreprocessingToLattigo());
 
     // Convert Alloc Ops to InPlace Ops
     pm.addPass(lattigo::createAllocToInPlace());
@@ -567,7 +575,7 @@ void torchLinalgToCkksBuilder(OpPassManager& manager,
   suboptions.ckksBootstrapWaterline = options.ckksBootstrapWaterline;
   suboptions.scalingModBits = options.scalingModBits;
   suboptions.firstModBits = options.firstModBits;
-  suboptions.splitPreprocessing = options.splitPreprocessing;
+  suboptions.enableSplitPreprocessing = options.enableSplitPreprocessing;
   suboptions.experimentalDisableLoopUnroll =
       options.experimentalDisableLoopUnroll;
   suboptions.usePublicKey = options.usePublicKey;