Orbit: Compression & SISO Partitioning

lib/Analysis/ILPBootstrapPlacementAnalysis/BUILD

@@ -7,15 +7,22 @@ package(
 
 cc_library(
     name = "ILPBootstrapPlacementAnalysis",
-    srcs = ["ILPBootstrapPlacementAnalysis.cpp"],
-    hdrs = ["ILPBootstrapPlacementAnalysis.h"],
+    srcs = [
+        "ILPBootstrapPlacementAnalysis.cpp",
+        "OpGrouping.cpp",
+    ],
+    hdrs = [
+        "ILPBootstrapPlacementAnalysis.h",
+        "OpGrouping.h",
+    ],
     deps = [
         "@com_google_absl//absl/status:statusor",
         "@com_google_ortools//ortools/math_opt/cpp:math_opt",
         "@com_google_ortools//ortools/math_opt/solvers:gscip_solver",
         "@heir//lib/Analysis/SecretnessAnalysis",
         "@heir//lib/Dialect/Mgmt/IR:Dialect",
         "@heir//lib/Dialect/Secret/IR:Dialect",
+        "@heir//lib/Dialect/TensorExt/IR:Dialect",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:Analysis",
         "@llvm-project//mlir:ArithDialect",

lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.h

@@ -14,10 +14,52 @@ class raw_ostream;
 
 namespace mlir {
 namespace heir {
+
+// A latency model of the form cost(level) = slope * level + intercept,
+// fitted from a per-level latency table.
+struct LinearCost {
+  double slope = 0.0;
+  double intercept = 0.0;
+};
+
+// Costs used by the ILP objective. Bootstrap and rescale management decisions
+// are charged constant costs. When hasLevelCosts is set, each tracked op is
+// additionally charged a level-dependent latency at its input level,
+// distinguishing ciphertext-ciphertext (CtCt) from ciphertext-plaintext (CtPt)
+// operands.
+struct OpCostModel {
+  double bootstrapCost = 0.0;
+  double rescaleCost = 0.0;
+  bool hasLevelCosts = false;
+  LinearCost addCtCt;
+  LinearCost addCtPt;
+  LinearCost mulCtCt;
+  LinearCost mulCtPt;
+  LinearCost rotate;
+  LinearCost negate;
+};
+
 class ILPBootstrapPlacementAnalysis {
  public:
   enum class ScaleMode { kCKKS, kLevelOnly };
 
+  struct Options {
+    int bootstrapWaterline = 0;
+    int scaleWaterline = 0;
+    int scaleFactorBits = 0;
+    int bootstrapLevelLowerBound = 0;
+    // Group structurally equivalent ops so they share ILP variables (Orbit's
+    // compression). Grouped and ungrouped models have
+    // the same optimum restricted to symmetric solutions.
+    bool compress = true;
+    // Minimum number of ops per SISO partition (Orbit's delta). Partitions
+    // are solved independently under enumerated boundary states and stitched
+    // by dynamic programming.
+    int partitionMinSize = 100;
+    OpCostModel costModel;
+    ScaleMode scaleMode = ScaleMode::kCKKS;
+  };
+
   struct NodeManagement {
     Value value;
     int inputLevel;
@@ -37,19 +79,8 @@ class ILPBootstrapPlacementAnalysis {
   };
 
   ILPBootstrapPlacementAnalysis(Operation* op, DataFlowSolver* solver,
-                                int bootstrapWaterline, int scaleWaterline,
-                                int scaleFactorBits,
-                                int bootstrapLevelLowerBound, int bootstrapCost,
-                                int rescaleCost, ScaleMode scaleMode)
-      : opToRunOn(op),
-        solver(solver),
-        bootstrapWaterline(bootstrapWaterline),
-        scaleWaterline(scaleWaterline),
-        scaleFactorBits(scaleFactorBits),
-        bootstrapLevelLowerBound(bootstrapLevelLowerBound),
-        bootstrapCost(bootstrapCost),
-        rescaleCost(rescaleCost),
-        scaleMode(scaleMode) {}
+                                const Options& options)
+      : opToRunOn(op), solver(solver), options(options) {}
   ~ILPBootstrapPlacementAnalysis() = default;
 
   LogicalResult solve();
@@ -82,13 +113,7 @@ class ILPBootstrapPlacementAnalysis {
  private:
   Operation* opToRunOn;
   DataFlowSolver* solver;
-  int bootstrapWaterline;
-  int scaleWaterline;
-  int scaleFactorBits;
-  int bootstrapLevelLowerBound;
-  int bootstrapCost;
-  int rescaleCost;
-  ScaleMode scaleMode;
+  Options options;
   llvm::DenseMap<Operation*, bool> solution;
   llvm::DenseMap<Value, int> solutionLevelBeforeBootstrap;
   llvm::DenseMap<Value, int> solutionLevelAfterBootstrap;