Improve waterline bootstrapping

This change replaces the single greedy waterline-bootstrapping pattern with an analysis pass that marks all the places bootstraps should be inserted before a secondary pass that actually mutates the IR.

The original pattern was failing to compile a number of larger programs involving loops: if the op-result of a loop that exhausts all levels was later used as the input to a ct-ct mul, the pattern would not fire and it would try to modreduce after (or before) a level 0 ciphetext. The new analysis pass properly inserts a bootstrap in situations like this.

After this change, the hotword convolutional model successfully compiles in ~70 seconds on my dev machine.

update: this change also required a small change to the scale analysis code, so that it initializes plaintexts to the default scale in the forward pass for multiplications (which should always happen at the default scale), and updated backward exit states to respect function return annotations when present. This hardening step was needed because the change to how bootstraps are inserted necessitated adding additional adjust_scale ops (e.g., a bootstrap just before a modreduce) which don't have forward propagation in scale analysis.
PiperOrigin-RevId: 953572669

Author: j2kun

lib/Analysis/LevelAnalysis/BUILD

@@ -7,8 +7,14 @@ package(
 
 cc_library(
     name = "LevelAnalysis",
-    srcs = ["LevelAnalysis.cpp"],
-    hdrs = ["LevelAnalysis.h"],
+    srcs = [
+        "BootstrapWaterlineAnalysis.cpp",
+        "LevelAnalysis.cpp",
+    ],
+    hdrs = [
+        "BootstrapWaterlineAnalysis.h",
+        "LevelAnalysis.h",
+    ],
     deps = [
         "@heir//lib/Analysis:Utils",
         "@heir//lib/Analysis/SecretnessAnalysis",
@@ -21,6 +27,7 @@ cc_library(
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:Analysis",
         "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:ControlFlowInterfaces",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Support",

lib/Analysis/LevelAnalysis/BootstrapWaterlineAnalysis.cpp

@@ -0,0 +1,106 @@
+#include "lib/Analysis/LevelAnalysis/BootstrapWaterlineAnalysis.h"
+
+#include <cassert>
+#include <functional>
+
+#include "lib/Analysis/LevelAnalysis/LevelAnalysis.h"
+#include "lib/Analysis/Utils.h"
+#include "lib/Dialect/HEIRInterfaces.h"
+#include "lib/Dialect/Mgmt/IR/MgmtOps.h"
+#include "llvm/include/llvm/Support/Debug.h"               // from @llvm-project
+#include "mlir/include/mlir/Analysis/DataFlowFramework.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/Interfaces/CallInterfaces.h"   // from @llvm-project
+#include "mlir/include/mlir/Support/LLVM.h"                // from @llvm-project
+
+#define DEBUG_TYPE "bootstrap-waterline-analysis"
+
+namespace mlir {
+namespace heir {
+
+LogicalResult BootstrapWaterlineAnalysis::visitOperation(
+    Operation* op, ArrayRef<const BootstrapWaterlineLattice*> operands,
+    ArrayRef<BootstrapWaterlineLattice*> results) {
+  auto propagate = [&](Value value, const BootstrapWaterlineState& state) {
+    auto* lattice = getLatticeElement(value);
+    ChangeResult changed = lattice->join(state);
+    propagateIfChanged(lattice, changed);
+  };
+
+  // 1. Extract LevelState from operands
+  SmallVector<LevelState> operandLevelStates;
+  for (auto* operand : operands) {
+    operandLevelStates.push_back(operand->getValue().getLevelState());
+  }
+
+  // 2. Compute prospective level
+  LevelState prospectiveLevel = deriveResultLevel(op, operandLevelStates);
+  if (levelBudget > 0 && prospectiveLevel.isInt() &&
+      prospectiveLevel.getInt() > levelBudget) {
+    prospectiveLevel = LevelState(Invalid{});
+  }
+
+  // 3. Determine if we need to reset/bootstrap
+  LevelState resultLevel;
+  bool resultNeedsBootstrap = false;
+
+  if (isa<ResetsLevelOpInterface>(op)) {
+    resultLevel = prospectiveLevel;
+    resultNeedsBootstrap = false;
+  } else {
+    // We wait as long as possible to bootstrap, meaning that the level will
+    // remain at the waterline (i.e., level zero) until it hits a level-reducing
+    // op, at which point we have to mark the _operand_ as needing a bootstrap.
+    // But since that operand has already been processed by the analysis, we
+    // mark the op result and then patch it up by the pass that uses this
+    // analysis.
+    bool exceedsWaterline =
+        prospectiveLevel.isInvalid() ||
+        (prospectiveLevel.isInt() && prospectiveLevel.getInt() > waterline);
+
+    resultNeedsBootstrap = exceedsWaterline;
+    if (exceedsWaterline) {
+      if (auto reduceOp = dyn_cast<ReducesLevelOpInterface>(op)) {
+        resultLevel = LevelState(reduceOp.getLevelsToDrop());
+      } else {
+        resultLevel = LevelState(0);
+      }
+    } else {
+      resultLevel = prospectiveLevel;
+    }
+  }
+
+  BootstrapWaterlineState resultState(resultLevel, resultNeedsBootstrap);
+
+  LLVM_DEBUG({
+    llvm::dbgs() << "BWAnalysis: " << op->getName() << " prospective=";
+    prospectiveLevel.print(llvm::dbgs());
+    llvm::dbgs() << " -> result=";
+    resultState.print(llvm::dbgs());
+    llvm::dbgs() << "\n";
+  });
+
+  for (auto result : op->getOpResults()) {
+    if (isa<mgmt::InitOp>(op) || isSecretInternal(op, result)) {
+      propagate(result, resultState);
+    }
+  }
+
+  return success();
+}
+
+void BootstrapWaterlineAnalysis::visitExternalCall(
+    CallOpInterface call,
+    ArrayRef<const BootstrapWaterlineLattice*> argumentLattices,
+    ArrayRef<BootstrapWaterlineLattice*> resultLattices) {
+  auto callback =
+      std::bind(&BootstrapWaterlineAnalysis::propagateIfChangedWrapper, this,
+                std::placeholders::_1, std::placeholders::_2);
+  ::mlir::heir::visitExternalCall<BootstrapWaterlineState,
+                                  BootstrapWaterlineLattice>(
+      call, argumentLattices, resultLattices, callback);
+}
+
+}  // namespace heir
+}  // namespace mlir

lib/Analysis/LevelAnalysis/BootstrapWaterlineAnalysis.h

@@ -0,0 +1,98 @@
+#ifndef LIB_ANALYSIS_LEVELANALYSIS_BOOTSTRAPWATERLINEANALYSIS_H_
+#define LIB_ANALYSIS_LEVELANALYSIS_BOOTSTRAPWATERLINEANALYSIS_H_
+
+#include "lib/Analysis/LevelAnalysis/LevelAnalysis.h"
+#include "lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.h"
+#include "mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h"  // from @llvm-project
+#include "mlir/include/mlir/Analysis/DataFlowFramework.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/Interfaces/CallInterfaces.h"   // from @llvm-project
+#include "mlir/include/mlir/Support/LLVM.h"                // from @llvm-project
+
+namespace mlir {
+namespace heir {
+
+class BootstrapWaterlineState {
+ public:
+  BootstrapWaterlineState() : levelState(Uninit{}), needsBootstrap(false) {}
+  BootstrapWaterlineState(LevelState levelState, bool needsBootstrap)
+      : levelState(levelState), needsBootstrap(needsBootstrap) {}
+
+  LevelState getLevelState() const { return levelState; }
+  bool getNeedsBootstrap() const { return needsBootstrap; }
+
+  bool operator==(const BootstrapWaterlineState& other) const {
+    return levelState == other.levelState &&
+           needsBootstrap == other.needsBootstrap;
+  }
+
+  static BootstrapWaterlineState join(const BootstrapWaterlineState& lhs,
+                                      const BootstrapWaterlineState& rhs) {
+    return BootstrapWaterlineState(
+        LevelState::join(lhs.levelState, rhs.levelState),
+        lhs.needsBootstrap || rhs.needsBootstrap);
+  }
+
+  void print(llvm::raw_ostream& os) const {
+    os << "BWState(";
+    levelState.print(os);
+    os << ", needsBootstrap=" << (needsBootstrap ? "true" : "false") << ")";
+  }
+
+  friend llvm::raw_ostream& operator<<(llvm::raw_ostream& os,
+                                       const BootstrapWaterlineState& state) {
+    state.print(os);
+    return os;
+  }
+
+ private:
+  LevelState levelState;
+  bool needsBootstrap;
+};
+
+class BootstrapWaterlineLattice
+    : public dataflow::Lattice<BootstrapWaterlineState> {
+ public:
+  using Lattice::Lattice;
+};
+
+class BootstrapWaterlineAnalysis
+    : public dataflow::SparseForwardDataFlowAnalysis<BootstrapWaterlineLattice>,
+      public SecretnessAnalysisDependent<BootstrapWaterlineAnalysis> {
+ public:
+  BootstrapWaterlineAnalysis(DataFlowSolver& solver, int waterline = 20,
+                             int levelBudget = 20)
+      : dataflow::SparseForwardDataFlowAnalysis<BootstrapWaterlineLattice>(
+            solver),
+        waterline(waterline),
+        levelBudget(levelBudget) {}
+  friend class SecretnessAnalysisDependent<BootstrapWaterlineAnalysis>;
+
+  void setToEntryState(BootstrapWaterlineLattice* lattice) override {
+    propagateIfChanged(
+        lattice, lattice->join(BootstrapWaterlineState(LevelState(0), false)));
+  }
+
+  LogicalResult visitOperation(
+      Operation* op, ArrayRef<const BootstrapWaterlineLattice*> operands,
+      ArrayRef<BootstrapWaterlineLattice*> results) override;
+
+  void visitExternalCall(
+      CallOpInterface call,
+      ArrayRef<const BootstrapWaterlineLattice*> argumentLattices,
+      ArrayRef<BootstrapWaterlineLattice*> resultLattices) override;
+
+  void propagateIfChangedWrapper(AnalysisState* state, ChangeResult changed) {
+    propagateIfChanged(state, changed);
+  }
+
+ private:
+  int waterline;
+  int levelBudget;
+};
+
+}  // namespace heir
+}  // namespace mlir
+
+#endif  // LIB_ANALYSIS_LEVELANALYSIS_BOOTSTRAPWATERLINEANALYSIS_H_

lib/Analysis/LevelAnalysis/LevelAnalysis.cpp

@@ -37,12 +37,12 @@ namespace heir {
 // LevelAnalysis (Forward)
 //===----------------------------------------------------------------------===//
 [[maybe_unused]] static void debugLog(StringRef opName,
-                                      ArrayRef<const LevelLattice*> operands,
+                                      ArrayRef<LevelState> operands,
                                       const LevelState& result) {
   LLVM_DEBUG({
     llvm::dbgs() << "transferForward: " << opName << "(";
-    for (auto* operand : operands) {
-      operand->getValue().print(llvm::dbgs());
+    for (const auto& operand : operands) {
+      operand.print(llvm::dbgs());
       llvm::dbgs() << ", ";
     }
     llvm::dbgs() << ") = ";
@@ -52,7 +52,7 @@ namespace heir {
 };
 
 LevelState transferForward(ReducesLevelOpInterface op,
-                           ArrayRef<const LevelLattice*> operands) {
+                           ArrayRef<LevelState> operands) {
   unsigned operandIdx = op.getOperandToReduce().getOperandNumber();
   LevelState result = std::visit(
       Overloaded{
@@ -63,13 +63,13 @@ LevelState transferForward(ReducesLevelOpInterface op,
             return LevelState(val + op.getLevelsToDrop());
           },
       },
-      operands[operandIdx]->getValue().get());
+      operands[operandIdx].get());
   LLVM_DEBUG(debugLog("ReduceLevelOpInterface", operands, result));
   return result;
 }
 
 LevelState transferForward(ReducesAllLevelsOpInterface op,
-                           ArrayRef<const LevelLattice*> operands) {
+                           ArrayRef<LevelState> operands) {
   LevelState result = std::visit(
       Overloaded{
           // MaxLevel -> MaxLevel should result in a no-op, so technically
@@ -79,13 +79,13 @@ LevelState transferForward(ReducesAllLevelsOpInterface op,
           [](Invalid) -> LevelState { return LevelState(Invalid{}); },
           [](int val) -> LevelState { return LevelState(MaxLevel{}); },
       },
-      operands[0]->getValue().get());
+      operands[0].get());
   LLVM_DEBUG(debugLog("ReduceAllLevelsOpInterface", operands, result));
   return result;
 }
 
 LevelState transferForward(ResetsLevelOpInterface op,
-                           ArrayRef<const LevelLattice*> operands) {
+                           ArrayRef<LevelState> operands) {
   unsigned operandIdx = op.getOperandToReset().getOperandNumber();
   LevelState result = std::visit(
       Overloaded{
@@ -94,13 +94,12 @@ LevelState transferForward(ResetsLevelOpInterface op,
           [](Invalid) -> LevelState { return LevelState(Invalid{}); },
           [](int val) -> LevelState { return LevelState(0); },
       },
-      operands[operandIdx]->getValue().get());
+      operands[operandIdx].get());
   LLVM_DEBUG(debugLog("ResetsLevelOpInterface", operands, result));
   return result;
 }
 
-LevelState deriveResultLevel(Operation* op,
-                             ArrayRef<const LevelLattice*> operands) {
+LevelState deriveResultLevel(Operation* op, ArrayRef<LevelState> operands) {
   return llvm::TypeSwitch<Operation*, LevelState>(op)
       .Case<ResetsLevelOpInterface>(
           [&](auto op) -> LevelState { return transferForward(op, operands); })
@@ -110,8 +109,8 @@ LevelState deriveResultLevel(Operation* op,
           [&](auto op) -> LevelState { return transferForward(op, operands); })
       .Default([&](auto* op) -> LevelState {
         LevelState result;
-        for (auto* operandState : operands) {
-          result = LevelState::join(result, operandState->getValue());
+        for (const auto& operand : operands) {
+          result = LevelState::join(result, operand);
         }
         LLVM_DEBUG(debugLog(op->getName().getStringRef(), operands, result));
         return result;
@@ -127,7 +126,11 @@ LogicalResult LevelAnalysis::visitOperation(
     propagateIfChanged(lattice, changed);
   };
 
-  LevelState resultLevel = deriveResultLevel(op, operands);
+  SmallVector<LevelState> operandStates;
+  for (auto* operand : operands) {
+    operandStates.push_back(operand->getValue());
+  }
+  LevelState resultLevel = deriveResultLevel(op, operandStates);
   if (resultLevel.isInt() && resultLevel.getInt() > levelBudget) {
     resultLevel = LevelState(Invalid{});
   }

lib/Analysis/LevelAnalysis/LevelAnalysis.h

@@ -218,8 +218,7 @@ class LevelAnalysis
   int levelBudget;
 };
 
-LevelState deriveResultLevel(Operation* op,
-                             ArrayRef<const LevelLattice*> operands);
+LevelState deriveResultLevel(Operation* op, ArrayRef<LevelState> operands);
 
 /// Backward Analyze the level of plaintext operands of ct-pt ops.
 ///
@@ -265,14 +264,12 @@ std::optional<int> getMaxLevel(Operation* root);
 // Get the maximum level of SSA values in the op, from the data flow solver.
 int getMaxLevel(Operation* top, DataFlowSolver* solver);
 
-LevelState transferForward(mgmt::ModReduceOp op,
-                           ArrayRef<const LevelLattice*> operands);
+LevelState transferForward(mgmt::ModReduceOp op, ArrayRef<LevelState> operands);
 LevelState transferForward(mgmt::LevelReduceOp op,
-                           ArrayRef<const LevelLattice*> operands);
+                           ArrayRef<LevelState> operands);
 LevelState transferForward(mgmt::LevelReduceMinOp op,
-                           ArrayRef<const LevelLattice*> operands);
-LevelState transferForward(mgmt::BootstrapOp op,
-                           ArrayRef<const LevelLattice*> operands);
+                           ArrayRef<LevelState> operands);
+LevelState transferForward(mgmt::BootstrapOp op, ArrayRef<LevelState> operands);
 
 }  // namespace heir
 }  // namespace mlir