google · copybara-service · Jul 31, 2026 · Jul 31, 2026
diff --git a/lib/Analysis/LevelAnalysis/LevelAnalysis.h b/lib/Analysis/LevelAnalysis/LevelAnalysis.h
@@ -41,6 +41,8 @@
 namespace mlir {
 namespace heir {
 
+constexpr int kDefaultLevelBudget = 40;
+
 // A sentinel for the maximum allowable level before it is determined exactly
 // what the max level is. In the semantics of this analysis, levels start from 0
 // (the initial level) and go up to some max level (determined by a pass that
@@ -193,9 +195,9 @@ class LevelAnalysis
     : public dataflow::SparseForwardDataFlowAnalysis<LevelLattice>,
       public SecretnessAnalysisDependent<LevelAnalysis> {
  public:
-  LevelAnalysis(DataFlowSolver& solver, int levelBudget = 40)
+  LevelAnalysis(DataFlowSolver& solver, int levelBudget = 0)
       : dataflow::SparseForwardDataFlowAnalysis<LevelLattice>(solver),
-        levelBudget(levelBudget) {}
+        levelBudget(levelBudget > 0 ? levelBudget : kDefaultLevelBudget) {}
   friend class SecretnessAnalysisDependent<LevelAnalysis>;
 
   void setToEntryState(LevelLattice* lattice) override {

diff --git a/lib/Dialect/Mgmt/Transforms/AnnotateMgmt.cpp b/lib/Dialect/Mgmt/Transforms/AnnotateMgmt.cpp
@@ -81,7 +81,7 @@ struct AnnotateMgmt : impl::AnnotateMgmtBase<AnnotateMgmt> {
     SymbolTableCollection symbolTable;
     dataflow::loadBaselineAnalyses(solver);
     solver.load<SecretnessAnalysis>();
-    solver.load<LevelAnalysis>();
+    solver.load<LevelAnalysis>(levelBudget);
     solver.load<LevelAnalysisBackward>(symbolTable);
     solver.load<DimensionAnalysis>();
     solver.load<DimensionAnalysisBackward>(symbolTable);

diff --git a/lib/Dialect/Mgmt/Transforms/Passes.td b/lib/Dialect/Mgmt/Transforms/Passes.td
@@ -13,6 +13,8 @@ def AnnotateMgmt : Pass<"annotate-mgmt"> {
   let options = [
     Option<"baseLevel", "base-level", "int",
            /*default=*/"0", "Level to start counting from (used by B/FV)">,
+    Option<"levelBudget", "level-budget", "int",
+           /*default=*/"0", "Maximum level budget for analysis (0 = default)">,
   ];
 }
 

diff --git a/lib/Transforms/Halo/PartialUnrollForLevelConsumption.cpp b/lib/Transforms/Halo/PartialUnrollForLevelConsumption.cpp
@@ -36,7 +36,7 @@ struct PartialUnrollForLevelConsumption
     DataFlowSolver solver;
     dataflow::loadBaselineAnalyses(solver);
     solver.load<SecretnessAnalysis>();
-    solver.load<LevelAnalysis>();
+    solver.load<LevelAnalysis>(forceMaxLevel);
 
     if (failed(solver.initializeAndRun(getOperation()))) {
       getOperation()->emitOpError() << "Failed to run the analysis.\n";

diff --git a/lib/Transforms/Halo/Passes.td b/lib/Transforms/Halo/Passes.td
@@ -102,6 +102,10 @@ def RegionBranchOpLevelInvariance : Pass<"region-branch-op-level-invariance"> {
     "mlir::scf::SCFDialect",
     "mlir::heir::mgmt::MgmtDialect",
   ];
+  let options = [
+    Option<"levelBudget", "level-budget", "int",
+           /*default=*/"0", "Maximum level budget for analysis (0 = default)">,
+  ];
 }
 
 #endif  // LIB_TRANSFORMS_HALO_PASSES_TD_
diff --git a/lib/Transforms/Halo/RegionBranchOpLevelInvariance.cpp b/lib/Transforms/Halo/RegionBranchOpLevelInvariance.cpp
@@ -31,7 +31,7 @@ struct RegionBranchOpLevelInvariance
     SymbolTableCollection symbolTable;
     dataflow::loadBaselineAnalyses(solver);
     solver.load<SecretnessAnalysis>();
-    solver.load<LevelAnalysis>();
+    solver.load<LevelAnalysis>(levelBudget);
     solver.load<LevelAnalysisBackward>(symbolTable);
 
     if (failed(solver.initializeAndRun(getOperation()))) {

diff --git a/lib/Transforms/SecretInsertMgmt/Passes.td b/lib/Transforms/SecretInsertMgmt/Passes.td
@@ -133,7 +133,7 @@ def AnnotateBootstrapWaterline : Pass<"annotate-bootstrap-waterline", "ModuleOp"
   let options = [
     Option<"bootstrapWaterline", "bootstrap-waterline", "int",
            /*default=*/"2", "Waterline for insert bootstrap op">,
-    Option<"levelBudget", "level-budget", "int", /*default=*/"40",
+    Option<"levelBudget", "level-budget", "int", /*default=*/"0",
            "An optional maximum level budget for the pipeline to assume">,
   ];
 }

diff --git a/lib/Transforms/SecretInsertMgmt/Pipeline.cpp b/lib/Transforms/SecretInsertMgmt/Pipeline.cpp
@@ -43,10 +43,10 @@ void runSolver(Operation* top, DataFlowSolver& solver) {
   }
 }
 
-void makeAndRunSolver(Operation* top, DataFlowSolver& solver) {
+void makeAndRunSolver(Operation* top, DataFlowSolver& solver, int levelBudget) {
   dataflow::loadBaselineAnalyses(solver);
   solver.load<SecretnessAnalysis>();
-  solver.load<LevelAnalysis>();
+  solver.load<LevelAnalysis>(levelBudget);
   solver.load<MulDepthAnalysis>();
   runSolver(top, solver);
 }
@@ -65,11 +65,11 @@ void makeAndRunSecretnessAndMulDepthSolver(Operation* top,
   runSolver(top, solver);
 }
 
-void makeAndRunSecretnessAndLevelSolver(Operation* top,
-                                        DataFlowSolver& solver) {
+void makeAndRunSecretnessAndLevelSolver(Operation* top, DataFlowSolver& solver,
+                                        int levelBudget) {
   dataflow::loadBaselineAnalyses(solver);
   solver.load<SecretnessAnalysis>();
-  solver.load<LevelAnalysis>();
+  solver.load<LevelAnalysis>(levelBudget);
   runSolver(top, solver);
 }
 
@@ -94,7 +94,7 @@ LogicalResult runInsertMgmtPipeline(Operation* top,
 
   // Run Level Analysis to check for convergence
   DataFlowSolver levelSolver;
-  makeAndRunSolver(top, levelSolver);
+  makeAndRunSolver(top, levelSolver, options.levelBudget);
 
   auto nonInvariantLoops = getNonInvariantLoops(top, &levelSolver);
 
@@ -106,7 +106,7 @@ LogicalResult runInsertMgmtPipeline(Operation* top,
     bootstrapLoopIterArgs(loop, &secretnessSolver);
 
     DataFlowSolver freshLevelSolver;
-    makeAndRunSolver(top, freshLevelSolver);
+    makeAndRunSolver(top, freshLevelSolver, options.levelBudget);
     unrollLoopForLevelUtilization(loop, &freshLevelSolver, options.levelBudget);
   }
 
@@ -122,18 +122,20 @@ LogicalResult runInsertMgmtPipeline(Operation* top,
 
   // An if statement must have each branch producing the same level as a result,
   // so the branch with the higher level must insert a level_reduce op.
-  adjustLevelsForRegionBranchOps(top);
+  adjustLevelsForRegionBranchOps(top, options.levelBudget);
   adjustScalesForRegionBranchOps(top, &idCounter);
 
   LDBG(2) << "Handling cross level ops";
-  handleCrossLevelOps(top, &idCounter, options.includeFloats);
+  handleCrossLevelOps(top, &idCounter, options.includeFloats,
+                      options.levelBudget);
 
   LDBG(2) << "Handling cross mul depth ops";
-  handleCrossMulDepthOps(top, &idCounter, options.includeFloats);
+  handleCrossMulDepthOps(top, &idCounter, options.includeFloats,
+                         options.levelBudget);
 
   // An if statement must have each branch producing the same level as a result,
   // so the branch with the higher level must insert a level_reduce op.
-  adjustLevelsForRegionBranchOps(top);
+  adjustLevelsForRegionBranchOps(top, options.levelBudget);
   return success();
 }
 
@@ -207,9 +209,10 @@ void insertRelinearizeAfterMult(Operation* top, bool includeFloats) {
   (void)walkAndApplyPatterns(top, std::move(patterns));
 }
 
-void handleCrossLevelOps(Operation* top, int* idCounter, bool includeFloats) {
+void handleCrossLevelOps(Operation* top, int* idCounter, bool includeFloats,
+                         int levelBudget) {
   DataFlowSolver solver;
-  makeAndRunSecretnessAndLevelSolver(top, solver);
+  makeAndRunSecretnessAndLevelSolver(top, solver, levelBudget);
   MLIRContext* ctx = top->getContext();
   RewritePatternSet patterns(ctx);
   patterns.add<MatchCrossLevel<arith::AddIOp>, MatchCrossLevel<arith::SubIOp>,
@@ -225,10 +228,10 @@ void handleCrossLevelOps(Operation* top, int* idCounter, bool includeFloats) {
 // this only happen for before-mul but not include-first-mul case
 // at the first level, a Value can be both mulResult or not mulResult
 // we should match their scale by adding one adjust scale op
-void handleCrossMulDepthOps(Operation* top, int* idCounter,
-                            bool includeFloats) {
+void handleCrossMulDepthOps(Operation* top, int* idCounter, bool includeFloats,
+                            int levelBudget) {
   DataFlowSolver solver;
-  makeAndRunSolver(top, solver);
+  makeAndRunSolver(top, solver, levelBudget);
   MLIRContext* ctx = top->getContext();
   RewritePatternSet patterns(ctx);
   patterns
@@ -392,11 +395,11 @@ SmallVector<Operation*> getNonInvariantLoops(Operation* top,
   return nonInvariantLoops;
 }
 
-void adjustLevelsForRegionBranchOps(Operation* top) {
+void adjustLevelsForRegionBranchOps(Operation* top, int levelBudget) {
   LDBG(2) << "Adjusting levels for region branching ops";
   MLIRContext* ctx = top->getContext();
   DataFlowSolver solver;
-  makeAndRunSecretnessAndLevelSolver(top, solver);
+  makeAndRunSecretnessAndLevelSolver(top, solver, levelBudget);
 
   RewritePatternSet patterns(ctx);
   patterns.add<RegionBranchOpLevelInvariancePattern>(ctx, &solver);

diff --git a/lib/Transforms/SecretInsertMgmt/Pipeline.h b/lib/Transforms/SecretInsertMgmt/Pipeline.h
@@ -34,13 +34,15 @@ void insertModReduceBeforeOrAfterMult(Operation* top, bool afterMul,
 
 void insertRelinearizeAfterMult(Operation* top, bool includeFloats);
 
-void adjustLevelsForRegionBranchOps(Operation* top);
+void adjustLevelsForRegionBranchOps(Operation* top, int levelBudget);
 
 void adjustScalesForRegionBranchOps(Operation* top, int* idCounter);
 
-void handleCrossLevelOps(Operation* top, int* idCounter, bool includeFloats);
+void handleCrossLevelOps(Operation* top, int* idCounter, bool includeFloats,
+                         int levelBudget);
 
-void handleCrossMulDepthOps(Operation* top, int* idCounter, bool includeFloats);
+void handleCrossMulDepthOps(Operation* top, int* idCounter, bool includeFloats,
+                            int levelBudget);
 
 void insertBootstrapWaterLine(Operation* top, int bootstrapWaterline,
                               int levelBudget, bool includeFloats,

diff --git a/lib/Transforms/SecretInsertMgmt/SecretInsertMgmtBFV.cpp b/lib/Transforms/SecretInsertMgmt/SecretInsertMgmtBFV.cpp
@@ -84,6 +84,7 @@ struct SecretInsertMgmtBFV
     pipeline.addPass(createCanonicalizerPass());
     mgmt::AnnotateMgmtOptions annotateMgmtOptions;
     annotateMgmtOptions.baseLevel = level;
+    annotateMgmtOptions.levelBudget = level;
     pipeline.addPass(mgmt::createAnnotateMgmt(annotateMgmtOptions));
     (void)runPipeline(pipeline, getOperation());
   }

diff --git a/lib/Transforms/SecretInsertMgmt/SecretInsertMgmtBGV.cpp b/lib/Transforms/SecretInsertMgmt/SecretInsertMgmtBGV.cpp
@@ -78,7 +78,9 @@ struct SecretInsertMgmtBGV
     OpPassManager pipeline("builtin.module");
     pipeline.addPass(createCanonicalizerPass());
     pipeline.addPass(createCSEPass());
-    pipeline.addPass(mgmt::createAnnotateMgmt());
+    mgmt::AnnotateMgmtOptions annotateOptions;
+    annotateOptions.levelBudget = levelBudget;
+    pipeline.addPass(mgmt::createAnnotateMgmt(annotateOptions));
     (void)runPipeline(pipeline, getOperation());
   }
 };

diff --git a/lib/Transforms/SecretInsertMgmt/SecretInsertMgmtCKKS.cpp b/lib/Transforms/SecretInsertMgmt/SecretInsertMgmtCKKS.cpp
@@ -54,7 +54,9 @@ struct SecretInsertMgmtCKKS
     OpPassManager pipeline("builtin.module");
     pipeline.addPass(createCanonicalizerPass());
     pipeline.addPass(createCSEPass());
-    pipeline.addPass(mgmt::createAnnotateMgmt());
+    mgmt::AnnotateMgmtOptions annotateOptions;
+    annotateOptions.levelBudget = levelBudget;
+    pipeline.addPass(mgmt::createAnnotateMgmt(annotateOptions));
     (void)runPipeline(pipeline, getOperation());
   }
 };

diff --git a/tests/Transforms/secret_insert_mgmt/ckks/slice.mlir b/tests/Transforms/secret_insert_mgmt/ckks/slice.mlir
@@ -0,0 +1,71 @@
+// RUN: heir-opt --secret-insert-mgmt-ckks=level-budget=41 %s | FileCheck %s
+
+// Ensure that bootstrapping is not applied to the secret tensor in the loop.
+
+module attributes {backend.lattigo, scheme.ckks} {
+  // CHECK: func.func @test_lenet_slice_loop
+  // CHECK: scf.for
+  // CHECK-NOT: mgmt.bootstrap
+  // CHECK: scf.yield
+  func.func @test_lenet_slice_loop(%ct_input: !secret.secret<tensor<1x8192xf32>>) -> !secret.secret<tensor<12x8192xf32>> {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %c12 = arith.constant 12 : index
+
+    %res = secret.generic(%ct_input: !secret.secret<tensor<1x8192xf32>>) {
+    ^body(%input: tensor<1x8192xf32>):
+      %empty = tensor.empty() : tensor<12x8192xf32>
+
+      %1 = arith.mulf %input, %input : tensor<1x8192xf32>
+      %2 = arith.mulf %1, %1 : tensor<1x8192xf32>
+      %3 = arith.mulf %2, %2 : tensor<1x8192xf32>
+      %4 = arith.mulf %3, %3 : tensor<1x8192xf32>
+      %5 = arith.mulf %4, %4 : tensor<1x8192xf32>
+      %6 = arith.mulf %5, %5 : tensor<1x8192xf32>
+      %7 = arith.mulf %6, %6 : tensor<1x8192xf32>
+      %8 = arith.mulf %7, %7 : tensor<1x8192xf32>
+      %9 = arith.mulf %8, %8 : tensor<1x8192xf32>
+      %10 = arith.mulf %9, %9 : tensor<1x8192xf32>
+      %11 = arith.mulf %10, %10 : tensor<1x8192xf32>
+      %12 = arith.mulf %11, %11 : tensor<1x8192xf32>
+      %13 = arith.mulf %12, %12 : tensor<1x8192xf32>
+      %14 = arith.mulf %13, %13 : tensor<1x8192xf32>
+      %15 = arith.mulf %14, %14 : tensor<1x8192xf32>
+      %16 = arith.mulf %15, %15 : tensor<1x8192xf32>
+      %17 = arith.mulf %16, %16 : tensor<1x8192xf32>
+      %18 = arith.mulf %17, %17 : tensor<1x8192xf32>
+      %19 = arith.mulf %18, %18 : tensor<1x8192xf32>
+      %20 = arith.mulf %19, %19 : tensor<1x8192xf32>
+      %21 = arith.mulf %20, %20 : tensor<1x8192xf32>
+      %22 = arith.mulf %21, %21 : tensor<1x8192xf32>
+      %23 = arith.mulf %22, %22 : tensor<1x8192xf32>
+      %24 = arith.mulf %23, %23 : tensor<1x8192xf32>
+      %25 = arith.mulf %24, %24 : tensor<1x8192xf32>
+      %26 = arith.mulf %25, %25 : tensor<1x8192xf32>
+      %27 = arith.mulf %26, %26 : tensor<1x8192xf32>
+      %28 = arith.mulf %27, %27 : tensor<1x8192xf32>
+      %29 = arith.mulf %28, %28 : tensor<1x8192xf32>
+      %30 = arith.mulf %29, %29 : tensor<1x8192xf32>
+      %31 = arith.mulf %30, %30 : tensor<1x8192xf32>
+      %32 = arith.mulf %31, %31 : tensor<1x8192xf32>
+      %33 = arith.mulf %32, %32 : tensor<1x8192xf32>
+      %34 = arith.mulf %33, %33 : tensor<1x8192xf32>
+      %35 = arith.mulf %34, %34 : tensor<1x8192xf32>
+      %36 = arith.mulf %35, %35 : tensor<1x8192xf32>
+      %37 = arith.mulf %36, %36 : tensor<1x8192xf32>
+      %38 = arith.mulf %37, %37 : tensor<1x8192xf32>
+      %39 = arith.mulf %38, %38 : tensor<1x8192xf32>
+      %40 = arith.mulf %39, %39 : tensor<1x8192xf32>
+      %41 = arith.mulf %40, %40 : tensor<1x8192xf32>
+
+      %loop = scf.for %idx = %c0 to %c12 step %c1 iter_args(%acc = %empty) -> (tensor<12x8192xf32>) {
+        %rot = tensor_ext.rotate %41, %idx : tensor<1x8192xf32>, index
+        %inserted = tensor.insert_slice %rot into %acc[%idx, 0] [1, 8192] [1, 1] : tensor<1x8192xf32> into tensor<12x8192xf32>
+        scf.yield %inserted : tensor<12x8192xf32>
+      }
+      secret.yield %loop : tensor<12x8192xf32>
+    } -> !secret.secret<tensor<12x8192xf32>>
+
+    return %res : !secret.secret<tensor<12x8192xf32>>
+  }
+}