From b72a6a18576cdf6ec1ecc7be6412e841487e4e11 Mon Sep 17 00:00:00 2001 From: Edward Chen Date: Mon, 13 Jul 2026 23:42:06 +0000 Subject: [PATCH 1/2] ilp-bootstrap-placement: level-dependent cost model and solver hygiene --- .../ILPBootstrapPlacementAnalysis/BUILD | 1 + .../ILPBootstrapPlacementAnalysis.cpp | 108 +++++++++++---- .../ILPBootstrapPlacementAnalysis.h | 36 ++++- .../ILPBootstrapPlacement.cpp | 131 +++++++++++++----- .../ILPBootstrapPlacement.td | 2 +- .../ILPBootstrapPlacement/README.md | 30 ++-- .../Transforms/ilp_bootstrap_placement/BUILD | 1 + .../orbit_cost_model.json | 47 ++++++- .../orbit_cost_model_error.mlir | 2 + .../orbit_incomplete_cost_model.json | 10 ++ .../orbit_level_dependent_costs.mlir | 52 +++++++ 11 files changed, 343 insertions(+), 77 deletions(-) create mode 100644 tests/Transforms/ilp_bootstrap_placement/orbit_incomplete_cost_model.json create mode 100644 tests/Transforms/ilp_bootstrap_placement/orbit_level_dependent_costs.mlir diff --git a/lib/Analysis/ILPBootstrapPlacementAnalysis/BUILD b/lib/Analysis/ILPBootstrapPlacementAnalysis/BUILD index b1c03782c8..bd7cb344cf 100644 --- a/lib/Analysis/ILPBootstrapPlacementAnalysis/BUILD +++ b/lib/Analysis/ILPBootstrapPlacementAnalysis/BUILD @@ -16,6 +16,7 @@ cc_library( "@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", diff --git a/lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.cpp b/lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.cpp index 4ba94e2006..07fdf30906 100644 --- a/lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.cpp +++ b/lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.cpp @@ -1,6 +1,7 @@ #include "lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.h" #include +#include #include #include #include @@ -9,6 +10,7 @@ #include "lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.h" #include "lib/Dialect/Mgmt/IR/MgmtAttributes.h" #include "lib/Dialect/Secret/IR/SecretOps.h" +#include "lib/Dialect/TensorExt/IR/TensorExtOps.h" #include "llvm/include/llvm/ADT/DenseMap.h" // from @llvm-project #include "llvm/include/llvm/ADT/STLExtras.h" // from @llvm-project #include "llvm/include/llvm/ADT/SmallVector.h" // from @llvm-project @@ -53,11 +55,11 @@ // plaintext constants contribute Sw // * node transitions relate each op's input state to each result state by // either direct rescale/modswitch management or a bootstrap transition -// * yielded result scales are constrained to explicit nonzero mgmt.mgmt -// scales -// on corresponding secret.generic results -// - Objective: minimize bootstrap and rescale costs, with a small tie-breaker -// favoring higher remaining levels. +// * a yielded value is pinned to the level annotated on its secret.generic +// result by an mgmt.mgmt attr, and to the annotated scale when nonzero +// - Objective: minimize total bootstrap and rescale cost. With a per-level cost +// model, also charge each tracked op its latency at its input level. A tiny +// per-level term breaks ties toward higher levels for output values. namespace math_opt = ::operations_research::math_opt; @@ -72,10 +74,29 @@ static bool isMultiplication(Operation* op) { return isa(op) || isa(op); } +static bool isAdditionLike(Operation* op) { + return isa(op); +} + static bool isConstantLike(Value value) { return value.getDefiningOp() != nullptr; } +// The level-dependent latency term for one op, or nullopt if the op class has +// no level-dependent cost. +static std::optional levelCostForOp(Operation* op, + DataFlowSolver* solver, + const OpCostModel& costModel) { + bool ctCt = llvm::count_if(op->getOperands(), [&](auto opd) { + return isSecret(opd, solver); + }) >= 2; + if (isMultiplication(op)) return ctCt ? costModel.mulCtCt : costModel.mulCtPt; + if (isAdditionLike(op)) return ctCt ? costModel.addCtCt : costModel.addCtPt; + if (isa(op)) return costModel.negate; + if (isa(op)) return costModel.rotate; + return std::nullopt; +} + static int roundedValue(const math_opt::VariableMap& varMap, const math_opt::Variable& var) { return static_cast(std::round(varMap.at(var))); @@ -331,10 +352,10 @@ static void addOperandEdgeConstraints(ILPModelState& state) { } } -// Add output-boundary scale constraints for values yielded from secret.generic. -// A yield is constrained only when the corresponding generic result has an -// explicit nonzero mgmt.mgmt scale; otherwise the ILP may choose any supported -// result scale and the later annotation pass records that chosen state. +// Add output-boundary constraints for values yielded from secret.generic. +// When the corresponding generic result carries an explicit mgmt.mgmt attr, +// the yielded value's level is pinned to the annotated level, and (in CKKS +// mode) a nonzero annotated scale pins the yielded value's scale. static LogicalResult addYieldConstraints(ILPModelState& state) { auto genericOp = cast(state.body->getParentOp()); for (Operation& op : state.body->getOperations()) { @@ -342,12 +363,26 @@ static LogicalResult addYieldConstraints(ILPModelState& state) { if (!yieldOp) continue; for (auto [index, operand] : llvm::enumerate(yieldOp->getOperands())) { if (!isSecret(operand, state.solver)) continue; - if (!state.valueScaleVars.contains(operand)) continue; - if (state.levelOnly) continue; + if (!state.valueLevelVars.contains(operand)) continue; mgmt::MgmtAttr mgmtAttr = mgmt::findMgmtAttrAssociatedWith(genericOp.getResult(index)); - if (!mgmtAttr || mgmtAttr.getScale() == 0) continue; + if (!mgmtAttr) continue; + + int resultLevel = mgmtAttr.getLevel(); + if (resultLevel < 0 || resultLevel > state.bootstrapWaterline) { + genericOp->emitError() + << "cannot constrain yielded value " << index + << " from secret.generic result mgmt.mgmt level " << resultLevel + << "; expected level in [0, " << state.bootstrapWaterline << "]"; + return failure(); + } + state.model.AddLinearConstraint( + state.valueLevelVars.at(operand) == resultLevel, + "yieldResultLevel" + std::to_string(index)); + + if (state.levelOnly || mgmtAttr.getScale() == 0) continue; + if (!state.valueScaleVars.contains(operand)) continue; int resultScale = mgmtAttr.getScale(); if (resultScale < state.sw || resultScale > state.scaleMax) { @@ -419,24 +454,33 @@ static void addNodeTransitionConstraints(ILPModelState& state, } } -static void addObjective(ILPModelState& state, int bootstrapCost, - int rescaleCost) { +static void addObjective(ILPModelState& state, const OpCostModel& costModel) { math_opt::LinearExpression objective; for (auto& [op, bootstrapVar] : state.bootstrapVars) { - objective += bootstrapCost * bootstrapVar; + objective += costModel.bootstrapCost * bootstrapVar; } for (auto& [op, rescaleVar] : state.nodeRescaleVars) { - objective += rescaleCost * rescaleVar; + objective += costModel.rescaleCost * rescaleVar; } for (auto& [operand, rescaleVar] : state.edgeRescaleVars) { - objective += rescaleCost * rescaleVar; + objective += costModel.rescaleCost * rescaleVar; } - // Tie-breaker: level constraints are one-sided, so among equal-cost - // solutions the solver could pick gratuitously low levels (free modswitches). - // The small negative weight prefers the highest feasible level for each - // value without outweighing a unit of bootstrap/rescale cost. TODO: remove - // in the next iteration, when the objective minimizes performance cost and - // per-level operation costs make level choices matter directly. + // Level-dependent op latency: each tracked op is charged + // slope * inputLevel + intercept for its cost class, so the solver prefers + // to run expensive ops (muls, rotations) at low levels. + if (costModel.hasLevelCosts) { + for (Operation* op : state.trackedOps) { + auto cost = levelCostForOp(op, state.solver, costModel); + if (!cost.has_value()) continue; + objective += cost->slope * state.inputLevelVars.at(op) + cost->intercept; + } + } + // Tie-breaker on value (result) levels: level constraints are one-sided, so + // among equal-cost solutions the solver could pick gratuitously low levels + // (free modswitches decoded as spurious level_reduce ops). The small + // negative weight prefers the highest feasible level for each value. Op + // *input* levels are separate variables and get real downward pressure from + // the level-dependent latency terms above, so the two do not conflict. for (auto& [value, levelVar] : state.valueLevelVars) { objective += -0.001 * levelVar; } @@ -497,7 +541,7 @@ LogicalResult ILPBootstrapPlacementAnalysis::solve() { addOperandEdgeConstraints(state); if (failed(addYieldConstraints(state))) return failure(); addNodeTransitionConstraints(state, bootstrapLevelLowerBound); - addObjective(state, bootstrapCost, rescaleCost); + addObjective(state, costModel); LLVM_DEBUG({ std::stringstream ss; @@ -505,8 +549,18 @@ LogicalResult ILPBootstrapPlacementAnalysis::solve() { llvm::dbgs() << "--- ILP model ---\n" << ss.str() << "--- end model ---\n"; }); + // Solve to a 1% relative optimality gap, matching Orbit's solver + // configuration (Gurobi MIPGap / CBC gapRel = 0.01). Proving full optimality + // often dominates solve time on large instances while improving the + // objective by less than measurement noise in the profiled cost models. On + // small instances the solver typically closes the gap entirely, so this + // rarely changes the chosen placement. + constexpr double kRelativeMipGap = 0.01; + math_opt::SolveArguments solveArgs; + solveArgs.parameters.relative_gap_tolerance = kRelativeMipGap; + const absl::StatusOr status = - math_opt::Solve(state.model, math_opt::SolverType::kGscip); + math_opt::Solve(state.model, math_opt::SolverType::kGscip, solveArgs); if (!status.ok()) { std::stringstream ss; ss << "Error solving the problem: " << status.status() << "\n"; @@ -520,8 +574,8 @@ LogicalResult ILPBootstrapPlacementAnalysis::solve() { case math_opt::TerminationReason::kFeasible: break; default: - llvm::errs() << "The problem does not have a feasible solution. " - "Termination status code: " + llvm::errs() << "No feasible solution found (the problem may be " + "infeasible). Termination status code: " << static_cast(result.termination.reason) << "\n"; return failure(); } diff --git a/lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.h b/lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.h index f9200613c4..ac5918a0e3 100644 --- a/lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.h +++ b/lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.h @@ -14,6 +14,31 @@ 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 }; @@ -39,16 +64,16 @@ class ILPBootstrapPlacementAnalysis { ILPBootstrapPlacementAnalysis(Operation* op, DataFlowSolver* solver, int bootstrapWaterline, int scaleWaterline, int scaleFactorBits, - int bootstrapLevelLowerBound, int bootstrapCost, - int rescaleCost, ScaleMode scaleMode) + int bootstrapLevelLowerBound, + const OpCostModel& costModel, + ScaleMode scaleMode) : opToRunOn(op), solver(solver), bootstrapWaterline(bootstrapWaterline), scaleWaterline(scaleWaterline), scaleFactorBits(scaleFactorBits), bootstrapLevelLowerBound(bootstrapLevelLowerBound), - bootstrapCost(bootstrapCost), - rescaleCost(rescaleCost), + costModel(costModel), scaleMode(scaleMode) {} ~ILPBootstrapPlacementAnalysis() = default; @@ -86,8 +111,7 @@ class ILPBootstrapPlacementAnalysis { int scaleWaterline; int scaleFactorBits; int bootstrapLevelLowerBound; - int bootstrapCost; - int rescaleCost; + OpCostModel costModel; ScaleMode scaleMode; llvm::DenseMap solution; llvm::DenseMap solutionLevelBeforeBootstrap; diff --git a/lib/Transforms/ILPBootstrapPlacement/ILPBootstrapPlacement.cpp b/lib/Transforms/ILPBootstrapPlacement/ILPBootstrapPlacement.cpp index 7152aa17f3..a3cce02936 100644 --- a/lib/Transforms/ILPBootstrapPlacement/ILPBootstrapPlacement.cpp +++ b/lib/Transforms/ILPBootstrapPlacement/ILPBootstrapPlacement.cpp @@ -1,6 +1,5 @@ #include "lib/Transforms/ILPBootstrapPlacement/ILPBootstrapPlacement.h" -#include #include #include #include @@ -11,6 +10,8 @@ #include "lib/Dialect/Mgmt/Transforms/AnnotateMgmt.h" #include "lib/Dialect/Secret/IR/SecretOps.h" #include "lib/Transforms/SecretInsertMgmt/Pipeline.h" +#include "llvm/include/llvm/ADT/STLExtras.h" // from @llvm-project +#include "llvm/include/llvm/ADT/SmallVector.h" // from @llvm-project #include "llvm/include/llvm/Support/Debug.h" // from @llvm-project #include "llvm/include/llvm/Support/JSON.h" // from @llvm-project #include "llvm/include/llvm/Support/MemoryBuffer.h" // from @llvm-project @@ -33,32 +34,73 @@ namespace heir { #define GEN_PASS_DEF_ILPBOOTSTRAPPLACEMENT #include "lib/Transforms/ILPBootstrapPlacement/ILPBootstrapPlacement.h.inc" -struct OrbitCostModel { - int bootstrapCost; - int rescaleCost; -}; - -static std::optional averagePositiveLatency(const llvm::json::Object& root, - llvm::StringRef opName) { - const llvm::json::Object* latencyTable = root.getObject("latencyTable"); - if (!latencyTable) return std::nullopt; - - const llvm::json::Array* latencies = latencyTable->getArray(opName); +// Read one op's per-level latency array (index i holds the latency at level +// i + 1). +static std::optional> readLatencies( + const llvm::json::Object& latencyTable, llvm::StringRef opName) { + const llvm::json::Array* latencies = latencyTable.getArray(opName); if (!latencies) return std::nullopt; - double sum = 0; - int count = 0; + SmallVector values; for (const llvm::json::Value& latencyValue : *latencies) { std::optional latency = latencyValue.getAsNumber(); - if (!latency || *latency <= 0) continue; - sum += *latency; + if (!latency) continue; + values.push_back(*latency); + } + if (values.empty()) return std::nullopt; + return values; +} + +static std::optional averagePositiveLatency( + const llvm::json::Object& latencyTable, llvm::StringRef opName) { + auto values = readLatencies(latencyTable, opName); + if (!values) return std::nullopt; + double sum = 0; + int count = 0; + for (double value : *values) { + if (value <= 0) continue; + sum += value; ++count; } if (count == 0) return std::nullopt; - return static_cast(std::llround(sum / count)); + return sum / count; } -static FailureOr loadOrbitCostModel(llvm::StringRef path) { +static std::optional maxPositiveLatency( + const llvm::json::Object& latencyTable, llvm::StringRef opName) { + auto values = readLatencies(latencyTable, opName); + if (!values) return std::nullopt; + double maxValue = *llvm::max_element(*values); + if (maxValue <= 0) return std::nullopt; + return maxValue; +} + +// Least-squares fit of cost(level) = slope * level + intercept over a +// per-level latency array, where array index i holds the latency at level +// i + 1. +static std::optional fitLinearCost( + const llvm::json::Object& latencyTable, llvm::StringRef opName) { + auto values = readLatencies(latencyTable, opName); + if (!values || values->empty()) return std::nullopt; + + int n = values->size(); + if (n == 1) return LinearCost{0.0, (*values)[0]}; + + double sumX = 0, sumY = 0, sumXY = 0, sumXX = 0; + for (int i = 0; i < n; ++i) { + double x = i + 1; + double y = (*values)[i]; + sumX += x; + sumY += y; + sumXY += x * y; + sumXX += x * x; + } + double slope = (n * sumXY - sumX * sumY) / (n * sumXX - sumX * sumX); + double intercept = (sumY - slope * sumX) / n; + return LinearCost{slope, intercept}; +} + +static FailureOr loadOrbitCostModel(llvm::StringRef path) { auto bufferOrError = llvm::MemoryBuffer::getFile(path); if (!bufferOrError) return failure(); @@ -71,14 +113,40 @@ static FailureOr loadOrbitCostModel(llvm::StringRef path) { const llvm::json::Object* root = parsed->getAsObject(); if (!root) return failure(); + const llvm::json::Object* latencyTable = root->getObject("latencyTable"); + if (!latencyTable) return failure(); + + // Bootstrap and rescale enter the objective as constant per-decision costs: + // bootstrap is the average of positive samples (levels below the bootstrap + // range are typically recorded as zero) and rescale is the per-level + // maximum. + std::optional bootstrapCost = + averagePositiveLatency(*latencyTable, "bootstrap"); + std::optional rescaleCost = + maxPositiveLatency(*latencyTable, "rescale"); + if (!bootstrapCost || !rescaleCost) return failure(); + + OpCostModel costModel; + costModel.bootstrapCost = *bootstrapCost; + costModel.rescaleCost = *rescaleCost; + + struct LevelCostKey { + llvm::StringRef name; + LinearCost* target; + }; + LevelCostKey levelCostKeys[] = { + {"addCtCt", &costModel.addCtCt}, {"addCtPt", &costModel.addCtPt}, + {"mulCtCt", &costModel.mulCtCt}, {"mulCtPt", &costModel.mulCtPt}, + {"rotate", &costModel.rotate}, {"negate", &costModel.negate}, + }; + for (auto& [key, target] : levelCostKeys) { + std::optional fitted = fitLinearCost(*latencyTable, key); + if (!fitted) return failure(); + *target = *fitted; + } + costModel.hasLevelCosts = true; - std::optional parsedBootstrapCost = - averagePositiveLatency(*root, "bootstrap"); - std::optional parsedRescaleCost = - averagePositiveLatency(*root, "rescale"); - if (!parsedBootstrapCost || !parsedRescaleCost) return failure(); - - return OrbitCostModel{*parsedBootstrapCost, *parsedRescaleCost}; + return costModel; } struct ILPBootstrapPlacement @@ -143,10 +211,11 @@ struct ILPBootstrapPlacement nodeManagement, SmallVector* edgeManagement) { - int effectiveBootstrapCost = bootstrapCost; - int effectiveRescaleCost = rescaleCost; + OpCostModel effectiveCostModel; + effectiveCostModel.bootstrapCost = bootstrapCost; + effectiveCostModel.rescaleCost = rescaleCost; if (!orbitCostModel.empty()) { - FailureOr loadedCostModel = + FailureOr loadedCostModel = loadOrbitCostModel(orbitCostModel); if (failed(loadedCostModel)) { llvm::errs() << "failed to load Orbit cost model from `" @@ -155,15 +224,13 @@ struct ILPBootstrapPlacement << orbitCostModel << "`"; return failure(); } - effectiveBootstrapCost = loadedCostModel->bootstrapCost; - effectiveRescaleCost = loadedCostModel->rescaleCost; + effectiveCostModel = *loadedCostModel; } ILPBootstrapPlacementAnalysis analysis( genericOp, solver, bootstrapWaterline, scaleConfig.scaleWaterline, scaleConfig.scaleFactorBits, bootstrapLevelLowerBound, - effectiveBootstrapCost, effectiveRescaleCost, - scaleConfig.analysisScaleMode()); + effectiveCostModel, scaleConfig.analysisScaleMode()); if (failed(analysis.solve())) { genericOp->emitError( "Failed to solve the bootstrap placement optimization problem"); diff --git a/lib/Transforms/ILPBootstrapPlacement/ILPBootstrapPlacement.td b/lib/Transforms/ILPBootstrapPlacement/ILPBootstrapPlacement.td index 83a5825a35..da3f327353 100644 --- a/lib/Transforms/ILPBootstrapPlacement/ILPBootstrapPlacement.td +++ b/lib/Transforms/ILPBootstrapPlacement/ILPBootstrapPlacement.td @@ -57,7 +57,7 @@ def ILPBootstrapPlacement : Pass <"ilp-bootstrap-placement"> { "orbit-cost-model", "std::string", /*default=*/"""", - "Path to a JSON cost model. When provided, bootstrap-cost and rescale-cost are loaded from latencyTable.bootstrap and latencyTable.rescale. See lib/Transforms/ILPBootstrapPlacement/README.md for the schema.">, + "Path to a JSON cost model. See lib/Transforms/ILPBootstrapPlacement/README.md for the schema.">, Option<"bootstrapCost", "bootstrap-cost", "int", diff --git a/lib/Transforms/ILPBootstrapPlacement/README.md b/lib/Transforms/ILPBootstrapPlacement/README.md index 07370c01fe..4e18150bbc 100644 --- a/lib/Transforms/ILPBootstrapPlacement/README.md +++ b/lib/Transforms/ILPBootstrapPlacement/README.md @@ -3,15 +3,29 @@ ## Cost Model JSON `--ilp-bootstrap-placement="orbit-cost-model=PATH"` loads an optional JSON cost -model and uses it to override the `bootstrap-cost` and `rescale-cost` pass -options. The units of latency are in microseconds. If level-dependent costs are -needed, the next step would be to extend this schema and the ILP objective. +model. The units of latency are in microseconds. -Required fields: +Every key under `latencyTable` maps to a per-level latency array, where index +`i` holds the latency at level `i + 1`. All keys are required; loading fails if +any is missing: -- `latencyTable.bootstrap`: numeric latency samples for one bootstrap chosen by - the ILP. -- `latencyTable.rescale`: numeric latency samples for one unit of rescale, - modreduce, or level-reduce management chosen by the ILP. +- `bootstrap`: the positive-sample average is the constant bootstrap cost in the + objective (levels below the bootstrappable range may be recorded as zero). + Overrides the `bootstrap-cost` option. +- `rescale`: the per-level maximum is the constant cost of one unit of rescale, + modreduce, or level-reduce management chosen by the ILP. Overrides the + `rescale-cost` option. +- `addCtCt`, `addCtPt`, `mulCtCt`, `mulCtPt`, `rotate`, `negate`: each array is + least-squares fitted to `cost(level) = slope * level + intercept`, and each + tracked op is charged its fitted cost at its ILP-chosen input level in the + objective. `CtCt` is the ciphertext-ciphertext variant of a binary op, `CtPt` + the ciphertext-plaintext variant. + +## Solver configuration + +The ILP is solved to a fixed 1% relative optimality gap with no time limit, +matching Orbit's solver configuration. Proving full optimality often dominates +solve time on large instances while improving the objective by less than +measurement noise in the profiled cost models. diff --git a/tests/Transforms/ilp_bootstrap_placement/BUILD b/tests/Transforms/ilp_bootstrap_placement/BUILD index b442021a1c..a5d4310d09 100644 --- a/tests/Transforms/ilp_bootstrap_placement/BUILD +++ b/tests/Transforms/ilp_bootstrap_placement/BUILD @@ -7,6 +7,7 @@ filegroup( srcs = [ "orbit_bad_cost_model.json", "orbit_cost_model.json", + "orbit_incomplete_cost_model.json", ], ) diff --git a/tests/Transforms/ilp_bootstrap_placement/orbit_cost_model.json b/tests/Transforms/ilp_bootstrap_placement/orbit_cost_model.json index 96eb2f686a..55d8930ff9 100644 --- a/tests/Transforms/ilp_bootstrap_placement/orbit_cost_model.json +++ b/tests/Transforms/ilp_bootstrap_placement/orbit_cost_model.json @@ -1,10 +1,51 @@ { "latencyTable": { - "bootstrap": [ - 69320650 + "addCtPt": [ + 50, + 100, + 150, + 200 + ], + "addCtCt": [ + 100, + 200, + 300, + 400 + ], + "mulCtPt": [ + 500, + 1000, + 1500, + 2000 + ], + "mulCtCt": [ + 1000, + 2000, + 3000, + 4000 + ], + "rotate": [ + 800, + 1600, + 2400, + 3200 + ], + "negate": [ + 20, + 40, + 60, + 80 ], "rescale": [ - 40988 + 100, + 100, + 100 + ], + "bootstrap": [ + 0, + 0, + 0, + 69000000 ] } } diff --git a/tests/Transforms/ilp_bootstrap_placement/orbit_cost_model_error.mlir b/tests/Transforms/ilp_bootstrap_placement/orbit_cost_model_error.mlir index 7b248a6ce7..87a273d67d 100644 --- a/tests/Transforms/ilp_bootstrap_placement/orbit_cost_model_error.mlir +++ b/tests/Transforms/ilp_bootstrap_placement/orbit_cost_model_error.mlir @@ -1,8 +1,10 @@ // RUN: not heir-opt --ilp-bootstrap-placement="orbit-cost-model=%S/does_not_exist.json" %s 2>&1 | FileCheck %s --check-prefix=MISSING // RUN: not heir-opt --ilp-bootstrap-placement="orbit-cost-model=%S/orbit_bad_cost_model.json" %s 2>&1 | FileCheck %s --check-prefix=MALFORMED +// RUN: not heir-opt --ilp-bootstrap-placement="orbit-cost-model=%S/orbit_incomplete_cost_model.json" %s 2>&1 | FileCheck %s --check-prefix=INCOMPLETE // MISSING: failed to load Orbit cost model // MALFORMED: failed to load Orbit cost model +// INCOMPLETE: failed to load Orbit cost model !pt_ty = tensor<8xf32> !ct_ty = !secret.secret diff --git a/tests/Transforms/ilp_bootstrap_placement/orbit_incomplete_cost_model.json b/tests/Transforms/ilp_bootstrap_placement/orbit_incomplete_cost_model.json new file mode 100644 index 0000000000..96eb2f686a --- /dev/null +++ b/tests/Transforms/ilp_bootstrap_placement/orbit_incomplete_cost_model.json @@ -0,0 +1,10 @@ +{ + "latencyTable": { + "bootstrap": [ + 69320650 + ], + "rescale": [ + 40988 + ] + } +} diff --git a/tests/Transforms/ilp_bootstrap_placement/orbit_level_dependent_costs.mlir b/tests/Transforms/ilp_bootstrap_placement/orbit_level_dependent_costs.mlir new file mode 100644 index 0000000000..36310c5b32 --- /dev/null +++ b/tests/Transforms/ilp_bootstrap_placement/orbit_level_dependent_costs.mlir @@ -0,0 +1,52 @@ +// RUN: heir-opt --ilp-bootstrap-placement="bootstrap-waterline=3 scale-waterline=51 scale-factor-bits=51" %s | FileCheck %s --check-prefix=CHECK-FLAT +// RUN: heir-opt --ilp-bootstrap-placement="bootstrap-waterline=3 scale-waterline=51 scale-factor-bits=51 orbit-cost-model=%S/orbit_cost_model.json" %s | FileCheck %s --check-prefix=CHECK-LEVEL + +// The circuit is mul -> adds -> mul, which needs exactly one rescale on each +// mul result to reach the annotated output state. The rescale between the two +// muls can be placed before or after the add chain without changing the +// rescale count, so the placement is decided by level preferences alone. +// +// With flat costs, the high-level tie-breaker keeps the adds at level 3 and +// rescales after the add chain. With Orbit-style level-dependent costs, adds +// are cheaper at lower levels, so the rescale moves directly after the first +// mul and the adds run at level 2. + +// CHECK-FLAT: func.func @level_dependent_rescale_timing +// CHECK-FLAT: arith.mulf %input0, %input0 +// CHECK-FLAT-NEXT: mgmt.relinearize +// CHECK-FLAT-NEXT: arith.addf +// CHECK-FLAT: mgmt.modreduce +// CHECK-FLAT: arith.mulf +// CHECK-FLAT-NEXT: mgmt.relinearize +// CHECK-FLAT-NEXT: mgmt.modreduce +// CHECK-FLAT-NOT: mgmt.bootstrap + +// CHECK-LEVEL: func.func @level_dependent_rescale_timing +// CHECK-LEVEL: arith.mulf %input0, %input0 +// CHECK-LEVEL-NEXT: mgmt.relinearize +// CHECK-LEVEL-NEXT: mgmt.modreduce +// CHECK-LEVEL-NEXT: arith.addf +// CHECK-LEVEL: arith.mulf +// CHECK-LEVEL-NEXT: mgmt.relinearize +// CHECK-LEVEL-NEXT: mgmt.modreduce +// CHECK-LEVEL-NOT: mgmt.bootstrap + +!pt_ty = tensor<8xf32> +!ct_ty = !secret.secret + +module attributes {scheme.ckks} { + func.func @level_dependent_rescale_timing( + %arg0: !ct_ty) + -> (!ct_ty {mgmt.mgmt = #mgmt.mgmt}) { + %0 = secret.generic(%arg0: !ct_ty) { + ^body(%input0: !pt_ty): + %m = arith.mulf %input0, %input0 : !pt_ty + %a1 = arith.addf %m, %m : !pt_ty + %a2 = arith.addf %a1, %a1 : !pt_ty + %a3 = arith.addf %a2, %a2 : !pt_ty + %mm = arith.mulf %a3, %a3 : !pt_ty + secret.yield %mm : !pt_ty + } -> (!ct_ty {mgmt.mgmt = #mgmt.mgmt}) + return %0 : !ct_ty + } +} From a610c631a423c7e690daf57b2d6f889fbe4b5b86 Mon Sep 17 00:00:00 2001 From: Edward Chen Date: Tue, 14 Jul 2026 00:42:09 +0000 Subject: [PATCH 2/2] ilp-bootstrap-placement: op grouping, compression, and SISO partitioning Port Orbit's three ILP-reduction techniques: addition-tree squashing and structural compression (auto_compress) merge equivalent ops into groups that share one set of ILP variables (exact, since op costs are linear in level), and single-input single-output partitioning solves large circuits piecewise under enumerated boundary states with a DP stitch. --- .../ILPBootstrapPlacementAnalysis/BUILD | 10 +- .../ILPBootstrapPlacementAnalysis.cpp | 986 +++++++++++++----- .../ILPBootstrapPlacementAnalysis.h | 39 +- .../OpGrouping.cpp | 344 ++++++ .../OpGrouping.h | 69 ++ .../ILPBootstrapPlacement.cpp | 18 +- .../ILPBootstrapPlacement.td | 10 + .../ILPBootstrapPlacement/README.md | 27 + .../orbit_compression.mlir | 35 + .../orbit_partition_ckks.mlir | 33 + .../orbit_siso_partition.mlir | 29 + 11 files changed, 1322 insertions(+), 278 deletions(-) create mode 100644 lib/Analysis/ILPBootstrapPlacementAnalysis/OpGrouping.cpp create mode 100644 lib/Analysis/ILPBootstrapPlacementAnalysis/OpGrouping.h create mode 100644 tests/Transforms/ilp_bootstrap_placement/orbit_compression.mlir create mode 100644 tests/Transforms/ilp_bootstrap_placement/orbit_partition_ckks.mlir create mode 100644 tests/Transforms/ilp_bootstrap_placement/orbit_siso_partition.mlir diff --git a/lib/Analysis/ILPBootstrapPlacementAnalysis/BUILD b/lib/Analysis/ILPBootstrapPlacementAnalysis/BUILD index bd7cb344cf..fb23c8afd6 100644 --- a/lib/Analysis/ILPBootstrapPlacementAnalysis/BUILD +++ b/lib/Analysis/ILPBootstrapPlacementAnalysis/BUILD @@ -7,8 +7,14 @@ 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", diff --git a/lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.cpp b/lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.cpp index 07fdf30906..8393a79862 100644 --- a/lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.cpp +++ b/lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.cpp @@ -1,17 +1,24 @@ #include "lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.h" +#include +#include #include +#include +#include #include #include #include #include +#include #include "absl/status/statusor.h" // from @com_google_absl +#include "lib/Analysis/ILPBootstrapPlacementAnalysis/OpGrouping.h" #include "lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.h" #include "lib/Dialect/Mgmt/IR/MgmtAttributes.h" #include "lib/Dialect/Secret/IR/SecretOps.h" #include "lib/Dialect/TensorExt/IR/TensorExtOps.h" #include "llvm/include/llvm/ADT/DenseMap.h" // from @llvm-project +#include "llvm/include/llvm/ADT/DenseSet.h" // from @llvm-project #include "llvm/include/llvm/ADT/STLExtras.h" // from @llvm-project #include "llvm/include/llvm/ADT/SmallVector.h" // from @llvm-project #include "llvm/include/llvm/Support/Debug.h" // from @llvm-project @@ -36,30 +43,39 @@ // keeping producer values, operand edges, and op results mutually feasible. // A later transform decodes the solution into mgmt operations. // -// ILP formulation: +// ILP formulation (per group of ops; see OpGrouping.h — grouped ops share one +// set of variables, and every objective term is scaled by the group's +// multiplicity): // - Variables: -// * level[value], scale[value] (in bits) for each secret SSA value -// * input_level[op], input_scale[op] for each tracked op -// * bootstrap[op], node_rescale[op] for management after an op -// * edge_rescale[use] and, for CKKS multiplication operands, edge_scale[use] -// for management before a consuming op +// * level[value], scale[value] (in bits) for each secret SSA value class +// * input_level[group], input_scale[group] for each tracked group +// * bootstrap[group], node_rescale[group] for management after a group +// * edge_rescale[edge] and, for CKKS multiplication operands, +// edge_scale[edge] for management before a consuming group, where an edge +// merges all operand uses of one producer class by one group // - Initialization: // * secret.generic body args are initialized from associated mgmt.mgmt attrs // when present, otherwise from (bootstrapWaterline, Sw) // - Constraints: // * bounds: levels are 0..bootstrapWaterline; CKKS scales are Sw..scaleMax // * level-only mode fixes all live scales to Sw -// * operand edges allow level/scale reduction before the consuming op +// * operand edges allow level/scale reduction before the consuming group // * non-multiplication operands share the consumer's input_scale // * CKKS multiplication input_scale is the sum of operand edge scales; // plaintext constants contribute Sw -// * node transitions relate each op's input state to each result state by +// * node transitions relate each group's input state to its result state by // either direct rescale/modswitch management or a bootstrap transition // * a yielded value is pinned to the level annotated on its secret.generic // result by an mgmt.mgmt attr, and to the annotated scale when nonzero // - Objective: minimize total bootstrap and rescale cost. With a per-level cost // model, also charge each tracked op its latency at its input level. A tiny // per-level term breaks ties toward higher levels for output values. +// +// Large circuits are additionally cut at single-input single-output (SISO) +// boundaries (Orbit's partitioning): each partition is solved independently +// under enumerated boundary (level, scale) states, producing a transfer table +// of boundary-in -> boundary-out costs, and a dynamic program stitches the +// per-partition solutions into a global placement. namespace math_opt = ::operations_research::math_opt; @@ -69,18 +85,9 @@ namespace mlir { namespace heir { using ScaleMode = ILPBootstrapPlacementAnalysis::ScaleMode; - -static bool isMultiplication(Operation* op) { - return isa(op) || isa(op); -} - -static bool isAdditionLike(Operation* op) { - return isa(op); -} - -static bool isConstantLike(Value value) { - return value.getDefiningOp() != nullptr; -} +using Options = ILPBootstrapPlacementAnalysis::Options; +using NodeManagement = ILPBootstrapPlacementAnalysis::NodeManagement; +using EdgeManagement = ILPBootstrapPlacementAnalysis::EdgeManagement; // The level-dependent latency term for one op, or nullopt if the op class has // no level-dependent cost. @@ -102,19 +109,73 @@ static int roundedValue(const math_opt::VariableMap& varMap, return static_cast(std::round(varMap.at(var))); } +// The smallest strictly positive cost that changes with a value's level: one +// rescale, or one level's worth of any op-latency slope. The value-level +// tie-breaker's total magnitude is capped below this so it can only break +// ties, never outweigh a real level-dependent cost. Intercepts and the +// bootstrap constant are excluded because they do not vary with a value +// level at the margin. +static double minLevelMarginalCost(const OpCostModel& costModel) { + double minCost = std::numeric_limits::infinity(); + auto consider = [&](double value) { + if (value > 0 && value < minCost) minCost = value; + }; + consider(costModel.rescaleCost); + if (costModel.hasLevelCosts) { + for (const LinearCost* cost : + {&costModel.addCtCt, &costModel.addCtPt, &costModel.mulCtCt, + &costModel.mulCtPt, &costModel.rotate, &costModel.negate}) { + consider(std::abs(cost->slope)); + } + } + return std::isinf(minCost) ? 1.0 : minCost; +} + +namespace { + +struct ValueState { + int level = 0; + int scale = 0; +}; + +// Boundary specification for one partition solve. +struct PartitionBoundary { + SmallVector, 4> pinnedInputs; + Value outputValue; + int outputLevel = 0; + // Yield pins apply only to the partition containing the yield. + bool applyYieldConstraints = true; + // Number of original ops in the partition, used to weigh the boundary + // output-scale pressure term. + int sizeInOps = 0; +}; + +// One merged operand edge: all uses of one producer value class by one +// consumer group share these decision variables. +struct GroupEdge { + math_opt::Variable rescaleVar; + math_opt::Variable scaleVar; + int weight = 0; +}; + struct ILPModelState { - ILPModelState(Block* body, DataFlowSolver* solver, int bootstrapWaterline, - int scaleWaterline, int scaleFactorBits, ScaleMode scaleMode) + ILPModelState(Block* body, DataFlowSolver* solver, const Options& options, + const OpGrouping& grouping, int groupBegin, int groupEnd, + const PartitionBoundary& boundary) : model("ILPBootstrapPlacementAnalysis"), body(body), solver(solver), - bootstrapWaterline(bootstrapWaterline), - levelOnly(scaleMode == ScaleMode::kLevelOnly), - sw(scaleWaterline), - sf(scaleFactorBits), + bootstrapWaterline(options.bootstrapWaterline), + levelOnly(options.scaleMode == ScaleMode::kLevelOnly), + sw(options.scaleWaterline), + sf(options.scaleFactorBits), scaleMax(levelOnly ? sw : sf + 2 * sw), bigM(bootstrapWaterline + 1), - scaleBigM(4 * scaleMax + sf * (bootstrapWaterline + 1)) {} + scaleBigM(4 * scaleMax + sf * (bootstrapWaterline + 1)), + grouping(grouping), + groupBegin(groupBegin), + groupEnd(groupEnd), + boundary(boundary) {} math_opt::Variable addLevelVar(int lower, int upper, const std::string& name) { @@ -140,6 +201,8 @@ struct ILPModelState { return ss.str(); } + Value canon(Value value) const { return grouping.canonicalValue(value); } + math_opt::Model model; Block* body; DataFlowSolver* solver; @@ -150,205 +213,242 @@ struct ILPModelState { int scaleMax; int bigM; int scaleBigM; + const OpGrouping& grouping; + int groupBegin; + int groupEnd; + const PartitionBoundary& boundary; int nextOpaqueId = 0; llvm::DenseMap opaqueIds; + // Value variables are keyed by canonical values (see OpGrouping::valueRep). llvm::DenseMap valueLevelVars; llvm::DenseMap valueScaleVars; + // Input variables are aliased to every member op of a group. llvm::DenseMap inputLevelVars; llvm::DenseMap inputScaleVars; + // Management variables are keyed by the group representative. llvm::DenseMap nodeRescaleVars; llvm::DenseMap bootstrapVars; + // Edge variables are aliased to every operand use merged into the edge. llvm::DenseMap edgeScaleVars; llvm::DenseMap edgeRescaleVars; + SmallVector edges; + llvm::DenseMap, int> edgeIndex; SmallVector trackedOps; }; -static LogicalResult addBodyArgumentVariables(ILPModelState& state) { - for (BlockArgument arg : state.body->getArguments()) { - if (!isSecret(arg, state.solver)) continue; +// One entry of the partition dynamic-programming transfer table: the decoded +// mgmt decisions and cost for one partition solved under one (input-state, +// output-level) boundary. The DP selects one of these per partition and the +// analysis adopts the chosen chain. +struct PartitionSolution { + double cost = 0; + ValueState outState; + std::pair inKey; + SmallVector nodeManagement; + SmallVector edgeManagement; + llvm::DenseMap bootstrapDecisions; + llvm::DenseMap levelBefore; + llvm::DenseMap levelAfter; +}; + +struct Partition { + int groupBegin; + int groupEnd; + // Canonical value live across the cut after this partition; null on the + // last partition. + Value cutValue; + int sizeInOps = 0; +}; - int initialLevel = state.bootstrapWaterline; - int initialScale = state.sw; +} // namespace + +// Determine the initial (level, scale) of each secret generic argument from +// associated mgmt.mgmt attrs, or the (bootstrapWaterline, Sw) defaults. +static LogicalResult computeArgInitStates( + Block* body, DataFlowSolver* solver, const Options& options, + SmallVector, 4>& pinnedInputs) { + bool levelOnly = options.scaleMode == ScaleMode::kLevelOnly; + int sw = options.scaleWaterline; + int scaleMax = levelOnly ? sw : options.scaleFactorBits + 2 * sw; + for (BlockArgument arg : body->getArguments()) { + if (!isSecret(arg, solver)) continue; + + int initialLevel = options.bootstrapWaterline; + int initialScale = sw; if (mgmt::MgmtAttr mgmtAttr = mgmt::findMgmtAttrAssociatedWith(arg)) { initialLevel = mgmtAttr.getLevel(); - if (!state.levelOnly && mgmtAttr.getScale() != 0) { + if (!levelOnly && mgmtAttr.getScale() != 0) { initialScale = mgmtAttr.getScale(); } } - Operation* parentOp = state.body->getParentOp(); - if (initialLevel < 0 || initialLevel > state.bootstrapWaterline) { + Operation* parentOp = body->getParentOp(); + if (initialLevel < 0 || initialLevel > options.bootstrapWaterline) { parentOp->emitError() << "cannot initialize ILP variable for secret.generic argument " << arg.getArgNumber() << " from mgmt.mgmt level " << initialLevel - << "; expected level in [0, " << state.bootstrapWaterline << "]"; + << "; expected level in [0, " << options.bootstrapWaterline << "]"; return failure(); } - if (initialScale < state.sw || initialScale > state.scaleMax) { + if (initialScale < sw || initialScale > scaleMax) { parentOp->emitError() << "cannot initialize ILP variable for secret.generic argument " << arg.getArgNumber() << " from mgmt.mgmt scale " << initialScale - << "; expected scale in [" << state.sw << ", " << state.scaleMax - << "]"; + << "; expected scale in [" << sw << ", " << scaleMax << "]"; return failure(); } - - std::stringstream ssLevel; - ssLevel << "levelArg" << arg.getArgNumber(); - auto levelVar = - state.addValueLevelVar(0, state.bootstrapWaterline, ssLevel.str()); - state.valueLevelVars.insert(std::make_pair(arg, levelVar)); - - std::stringstream ssScale; - ssScale << "scaleArg" << arg.getArgNumber(); - auto scaleVar = state.addScaleVar(ssScale.str()); - state.valueScaleVars.insert(std::make_pair(arg, scaleVar)); - - state.model.AddLinearConstraint( - levelVar == initialLevel, - "initLevelArg" + std::to_string(arg.getArgNumber())); - state.model.AddLinearConstraint( - scaleVar == initialScale, - "initScaleArg" + std::to_string(arg.getArgNumber())); + pinnedInputs.push_back({arg, {initialLevel, initialScale}}); } return success(); } -static bool hasSecretResult(Operation& op, DataFlowSolver* solver) { - return llvm::any_of(op.getResults(), [&](OpResult result) { - return isSecret(result, solver); - }); -} - -static bool shouldTrackOperation(Operation& op, DataFlowSolver* solver) { - return !isa(op) && hasSecretResult(op, solver); -} - -static void addOperationDecisionVariables(ILPModelState& state, Operation* op, - const std::string& opName) { - auto inputLevelVar = - state.addLevelVar(0, state.bootstrapWaterline, "inputLevel" + opName); - state.inputLevelVars.insert(std::make_pair(op, inputLevelVar)); - auto inputScaleVar = state.addScaleVar("inputScale" + opName); - state.inputScaleVars.insert(std::make_pair(op, inputScaleVar)); - auto nodeRescaleVar = - state.addLevelVar(0, state.bootstrapWaterline, "nodeRescale" + opName); - state.nodeRescaleVars.insert(std::make_pair(op, nodeRescaleVar)); - auto bootstrapVar = state.model.AddBinaryVariable("bootstrap" + opName); - state.bootstrapVars.insert(std::make_pair(op, bootstrapVar)); - - if (state.levelOnly) { - state.model.AddLinearConstraint(inputScaleVar == state.sw, - "levelOnlyInputScale" + opName); +static void addPinnedInputVariables(ILPModelState& state) { + int index = 0; + for (auto& [value, valueState] : state.boundary.pinnedInputs) { + std::string name = "Pinned" + std::to_string(index++); + auto levelVar = + state.addValueLevelVar(0, state.bootstrapWaterline, "level" + name); + state.valueLevelVars.insert(std::make_pair(value, levelVar)); + auto scaleVar = state.addScaleVar("scale" + name); + state.valueScaleVars.insert(std::make_pair(value, scaleVar)); + + state.model.AddLinearConstraint(levelVar == valueState.level, + "initLevel" + name); + state.model.AddLinearConstraint(scaleVar == valueState.scale, + "initScale" + name); } } -static void addOperationResultVariables(ILPModelState& state, Operation* op, - const std::string& opName) { - for (OpResult result : op->getResults()) { - if (!isSecret(result, state.solver)) continue; - - std::stringstream ssLevel; - ssLevel << "level" << opName << result.getResultNumber(); - auto levelVar = - state.addValueLevelVar(0, state.bootstrapWaterline, ssLevel.str()); - state.valueLevelVars.insert(std::make_pair(result, levelVar)); +static void addTrackedGroupVariables(ILPModelState& state) { + for (int gi = state.groupBegin; gi < state.groupEnd; ++gi) { + const OpGroup& group = state.grouping.groups[gi]; + Operation* rep = group.representative; + std::string opName = state.uniqueName(rep); + + auto inputLevelVar = + state.addLevelVar(0, state.bootstrapWaterline, "inputLevel" + opName); + auto inputScaleVar = state.addScaleVar("inputScale" + opName); + auto nodeRescaleVar = + state.addLevelVar(0, state.bootstrapWaterline, "nodeRescale" + opName); + auto bootstrapVar = state.model.AddBinaryVariable("bootstrap" + opName); + state.nodeRescaleVars.insert(std::make_pair(rep, nodeRescaleVar)); + state.bootstrapVars.insert(std::make_pair(rep, bootstrapVar)); + for (Operation* member : group.members) { + state.inputLevelVars.insert(std::make_pair(member, inputLevelVar)); + state.inputScaleVars.insert(std::make_pair(member, inputScaleVar)); + state.trackedOps.push_back(member); + } - std::stringstream ssScale; - ssScale << "scale" << opName << result.getResultNumber(); - auto scaleVar = state.addScaleVar(ssScale.str()); - state.valueScaleVars.insert(std::make_pair(result, scaleVar)); if (state.levelOnly) { - state.model.AddLinearConstraint( - scaleVar == state.sw, "levelOnlyResultScale" + opName + - std::to_string(result.getResultNumber())); + state.model.AddLinearConstraint(inputScaleVar == state.sw, + "levelOnlyInputScale" + opName); } - } -} - -static void addTrackedOperationVariables(ILPModelState& state) { - for (Operation& op : state.body->getOperations()) { - if (!shouldTrackOperation(op, state.solver)) continue; - state.trackedOps.push_back(&op); - std::string opName = state.uniqueName(&op); - addOperationDecisionVariables(state, &op, opName); - addOperationResultVariables(state, &op, opName); - } -} - -// Add constraints for one producer value flowing into one operand use of an op. -// The edge rescale variable models management inserted before the consumer. -// Non-multiplication ops share one input scale across all operands; CKKS -// multiplications use per-edge scales so their combined scale can be modeled -// separately in addMultiplicationInputScaleConstraint. -static void addSingleOperandEdgeConstraints(ILPModelState& state, Operation* op, - OpOperand& operandUse, - const std::string& opName) { - Value operand = operandUse.get(); - if (!isSecret(operand, state.solver)) return; - if (!state.valueLevelVars.contains(operand)) return; - - std::stringstream ss; - ss << opName << "Op" << operandUse.getOperandNumber(); - std::string edgeName = ss.str(); - - auto inputLevelVar = state.inputLevelVars.at(op); - auto inputScaleVar = state.inputScaleVars.at(op); - auto edgeRescaleVar = - state.addLevelVar(0, state.bootstrapWaterline, "edgeRescale" + edgeName); - state.edgeRescaleVars.insert(std::make_pair(&operandUse, edgeRescaleVar)); - - math_opt::Variable edgeScaleVar = - (!state.levelOnly && isMultiplication(op)) - ? state.addScaleVar("edgeScale" + edgeName) - : inputScaleVar; - state.edgeScaleVars.insert(std::make_pair(&operandUse, edgeScaleVar)); + for (OpResult result : rep->getResults()) { + if (!isSecret(result, state.solver)) continue; - state.model.AddLinearConstraint( - inputLevelVar <= state.valueLevelVars.at(operand) - edgeRescaleVar, - "edgeLevel" + edgeName); - state.model.AddLinearConstraint( - edgeScaleVar >= - state.valueScaleVars.at(operand) - state.sf * edgeRescaleVar, - "edgeScale" + edgeName); - if (state.levelOnly) { - state.model.AddLinearConstraint(edgeScaleVar == state.sw, - "levelOnlyEdgeScale" + edgeName); + std::stringstream ssLevel; + ssLevel << "level" << opName << result.getResultNumber(); + auto levelVar = + state.addValueLevelVar(0, state.bootstrapWaterline, ssLevel.str()); + state.valueLevelVars.insert(std::make_pair(result, levelVar)); + + std::stringstream ssScale; + ssScale << "scale" << opName << result.getResultNumber(); + auto scaleVar = state.addScaleVar(ssScale.str()); + state.valueScaleVars.insert(std::make_pair(result, scaleVar)); + if (state.levelOnly) { + state.model.AddLinearConstraint( + scaleVar == state.sw, "levelOnlyResultScale" + opName + + std::to_string(result.getResultNumber())); + } + } } } -// Add the CKKS multiplication scale-composition constraint. After each operand -// edge chooses its aligned scale, the multiplication's raw input scale is the -// sum of those operand scales; plaintext constants contribute the waterline -// scale. Result/output scale constraints are added by -// addNodeTransitionConstraints. -static void addMultiplicationInputScaleConstraint(ILPModelState& state, - Operation* op, - const std::string& opName) { - if (state.levelOnly || !isMultiplication(op)) return; - - math_opt::LinearExpression inputScale; - for (OpOperand& operandUse : op->getOpOperands()) { - Value operand = operandUse.get(); - if (isSecret(operand, state.solver) && - state.edgeScaleVars.contains(&operandUse)) { - inputScale += state.edgeScaleVars.at(&operandUse); - } else if (isConstantLike(operand)) { - inputScale += state.sw; +// Add constraints for the operand edges of each group. All uses of one +// producer value class by one group share a single edge: the edge rescale +// variable models management inserted before the consumers, charged once per +// original use in the objective. Non-multiplication groups share one input +// scale across all operands; CKKS multiplications use per-edge scales so +// their combined scale can be modeled separately below. +static void addGroupEdgeConstraints(ILPModelState& state) { + for (int gi = state.groupBegin; gi < state.groupEnd; ++gi) { + const OpGroup& group = state.grouping.groups[gi]; + Operation* rep = group.representative; + std::string opName = state.uniqueName(rep); + auto inputLevelVar = state.inputLevelVars.at(rep); + auto inputScaleVar = state.inputScaleVars.at(rep); + + for (Operation* member : group.members) { + // Weight counts distinct (producer class, consumer op) connections, so + // repeated uses of one producer by one op (e.g. squaring x*x) are one + // edge, matching Orbit's DiGraph edge model. Every operand use still + // registers the shared edge variables so the CKKS scale-composition sum + // and the decoder can see each slot. + DenseSet countedForMember; + for (OpOperand& operandUse : member->getOpOperands()) { + Value operand = operandUse.get(); + if (!isSecret(operand, state.solver)) continue; + if (Operation* def = operand.getDefiningOp()) { + auto it = state.grouping.groupIdOf.find(def); + if (it != state.grouping.groupIdOf.end() && it->second == gi) + continue; // interior edge of an addition tree + } + Value key = state.canon(operand); + if (!state.valueLevelVars.contains(key)) continue; + + auto [entry, inserted] = state.edgeIndex.try_emplace( + std::make_pair(key, gi), state.edges.size()); + if (inserted) { + std::string edgeName = + opName + "Edge" + std::to_string(state.edges.size()); + auto edgeRescaleVar = state.addLevelVar(0, state.bootstrapWaterline, + "edgeRescale" + edgeName); + math_opt::Variable edgeScaleVar = + (!state.levelOnly && group.isMultiplication) + ? state.addScaleVar("edgeScale" + edgeName) + : inputScaleVar; + + state.model.AddLinearConstraint( + inputLevelVar <= state.valueLevelVars.at(key) - edgeRescaleVar, + "edgeLevel" + edgeName); + state.model.AddLinearConstraint( + edgeScaleVar >= + state.valueScaleVars.at(key) - state.sf * edgeRescaleVar, + "edgeScale" + edgeName); + if (state.levelOnly) { + state.model.AddLinearConstraint(edgeScaleVar == state.sw, + "levelOnlyEdgeScale" + edgeName); + } + state.edges.push_back({edgeRescaleVar, edgeScaleVar, 0}); + } + GroupEdge& edge = state.edges[entry->second]; + if (countedForMember.insert(key).second) edge.weight += 1; + state.edgeRescaleVars.insert( + std::make_pair(&operandUse, edge.rescaleVar)); + state.edgeScaleVars.insert(std::make_pair(&operandUse, edge.scaleVar)); + } } - } - state.model.AddLinearConstraint(state.inputScaleVars.at(op) == inputScale, - "mulInputScale" + opName); -} -static void addOperandEdgeConstraints(ILPModelState& state) { - for (Operation* op : state.trackedOps) { - std::string opName = state.uniqueName(op); - for (OpOperand& operandUse : op->getOpOperands()) { - addSingleOperandEdgeConstraints(state, op, operandUse, opName); + // The CKKS multiplication scale-composition constraint: the group's raw + // input scale is the sum of the representative's operand edge scales; + // plaintext constants contribute the waterline scale. Any merged member + // satisfies the same constraint by symmetry. + if (!state.levelOnly && group.isMultiplication) { + math_opt::LinearExpression inputScale; + for (OpOperand& operandUse : rep->getOpOperands()) { + Value operand = operandUse.get(); + if (isSecret(operand, state.solver) && + state.edgeScaleVars.contains(&operandUse)) { + inputScale += state.edgeScaleVars.at(&operandUse); + } else if (isConstantLike(operand)) { + inputScale += state.sw; + } + } + state.model.AddLinearConstraint(inputScaleVar == inputScale, + "mulInputScale" + opName); } - addMultiplicationInputScaleConstraint(state, op, opName); } } @@ -357,13 +457,15 @@ static void addOperandEdgeConstraints(ILPModelState& state) { // the yielded value's level is pinned to the annotated level, and (in CKKS // mode) a nonzero annotated scale pins the yielded value's scale. static LogicalResult addYieldConstraints(ILPModelState& state) { + if (!state.boundary.applyYieldConstraints) return success(); auto genericOp = cast(state.body->getParentOp()); for (Operation& op : state.body->getOperations()) { auto yieldOp = dyn_cast(op); if (!yieldOp) continue; for (auto [index, operand] : llvm::enumerate(yieldOp->getOperands())) { if (!isSecret(operand, state.solver)) continue; - if (!state.valueLevelVars.contains(operand)) continue; + Value key = state.canon(operand); + if (!state.valueLevelVars.contains(key)) continue; mgmt::MgmtAttr mgmtAttr = mgmt::findMgmtAttrAssociatedWith(genericOp.getResult(index)); @@ -378,11 +480,11 @@ static LogicalResult addYieldConstraints(ILPModelState& state) { return failure(); } state.model.AddLinearConstraint( - state.valueLevelVars.at(operand) == resultLevel, + state.valueLevelVars.at(key) == resultLevel, "yieldResultLevel" + std::to_string(index)); if (state.levelOnly || mgmtAttr.getScale() == 0) continue; - if (!state.valueScaleVars.contains(operand)) continue; + if (!state.valueScaleVars.contains(key)) continue; int resultScale = mgmtAttr.getScale(); if (resultScale < state.sw || resultScale > state.scaleMax) { @@ -393,29 +495,49 @@ static LogicalResult addYieldConstraints(ILPModelState& state) { return failure(); } state.model.AddLinearConstraint( - state.valueScaleVars.at(operand) == resultScale, + state.valueScaleVars.at(key) == resultScale, "yieldResultScale" + std::to_string(index)); } } return success(); } -// Add constraints for each tracked op's result state after the op and any +// Pin the partition's boundary output level during boundary enumeration. +static void addOutputPinConstraints(ILPModelState& state) { + if (!state.boundary.outputValue) return; + Value key = state.canon(state.boundary.outputValue); + state.model.AddLinearConstraint( + state.valueLevelVars.at(key) == state.boundary.outputLevel, + "boundaryOutLevel"); + if (!state.levelOnly) { + // Backend feasibility at a partition boundary, matching Orbit's output + // scale/level relation for Lattigo: scale <= Sf * (level + 1) - margin. + constexpr int kBoundaryScaleMargin = 7; + state.model.AddLinearConstraint( + state.valueScaleVars.at(key) <= + state.sf * (state.boundary.outputLevel + 1) - kBoundaryScaleMargin, + "boundaryOutScale"); + } +} + +// Add constraints for each group's result state after the group and any // management chosen on the node. This is the result/output counterpart to the -// edge constraints: it relates the op's input level/scale to each secret result +// edge constraints: it relates the group's input level/scale to its result // through either a direct transition or a bootstrap transition. static void addNodeTransitionConstraints(ILPModelState& state, int bootstrapLevelLowerBound) { - for (Operation* op : state.trackedOps) { - std::string opName = state.uniqueName(op); - auto inputLevelVar = state.inputLevelVars.at(op); - auto inputScaleVar = state.inputScaleVars.at(op); - auto nodeRescaleVar = state.nodeRescaleVars.at(op); - auto bootstrapVar = state.bootstrapVars.at(op); + for (int gi = state.groupBegin; gi < state.groupEnd; ++gi) { + const OpGroup& group = state.grouping.groups[gi]; + Operation* rep = group.representative; + std::string opName = state.uniqueName(rep); + auto inputLevelVar = state.inputLevelVars.at(rep); + auto inputScaleVar = state.inputScaleVars.at(rep); + auto nodeRescaleVar = state.nodeRescaleVars.at(rep); + auto bootstrapVar = state.bootstrapVars.at(rep); int intrinsicLevelConsumption = - state.levelOnly && isMultiplication(op) ? 1 : 0; + state.levelOnly && group.isMultiplication ? 1 : 0; - for (OpResult result : op->getResults()) { + for (OpResult result : rep->getResults()) { if (!isSecret(result, state.solver)) continue; if (!state.valueLevelVars.contains(result)) continue; @@ -456,18 +578,22 @@ static void addNodeTransitionConstraints(ILPModelState& state, static void addObjective(ILPModelState& state, const OpCostModel& costModel) { math_opt::LinearExpression objective; - for (auto& [op, bootstrapVar] : state.bootstrapVars) { - objective += costModel.bootstrapCost * bootstrapVar; + // Management decisions are charged once per merged management site. + for (int gi = state.groupBegin; gi < state.groupEnd; ++gi) { + const OpGroup& group = state.grouping.groups[gi]; + Operation* rep = group.representative; + objective += + group.weight * costModel.bootstrapCost * state.bootstrapVars.at(rep); + objective += + group.weight * costModel.rescaleCost * state.nodeRescaleVars.at(rep); } - for (auto& [op, rescaleVar] : state.nodeRescaleVars) { - objective += costModel.rescaleCost * rescaleVar; - } - for (auto& [operand, rescaleVar] : state.edgeRescaleVars) { - objective += costModel.rescaleCost * rescaleVar; + for (const GroupEdge& edge : state.edges) { + objective += edge.weight * costModel.rescaleCost * edge.rescaleVar; } // Level-dependent op latency: each tracked op is charged // slope * inputLevel + intercept for its cost class, so the solver prefers - // to run expensive ops (muls, rotations) at low levels. + // to run expensive ops (muls, rotations) at low levels. Members of a group + // share one input level variable, so this sums member costs at that level. if (costModel.hasLevelCosts) { for (Operation* op : state.trackedOps) { auto cost = levelCostForOp(op, state.solver, costModel); @@ -477,71 +603,145 @@ static void addObjective(ILPModelState& state, const OpCostModel& costModel) { } // Tie-breaker on value (result) levels: level constraints are one-sided, so // among equal-cost solutions the solver could pick gratuitously low levels - // (free modswitches decoded as spurious level_reduce ops). The small - // negative weight prefers the highest feasible level for each value. Op - // *input* levels are separate variables and get real downward pressure from - // the level-dependent latency terms above, so the two do not conflict. + // (free modswitches decoded as spurious level_reduce ops). A small negative + // weight prefers the highest feasible level for each value. Op *input* + // levels are separate variables and get real downward pressure from the + // level-dependent latency terms above, so the two do not conflict. + // + // The per-value weight is nominally kEpsilon, but its total magnitude + // (summed over every value level, each at most bootstrapWaterline) is capped + // strictly below the smallest real level-marginal cost, so on large regions + // it can only break ties and never outweigh a genuine cost difference. Small + // regions keep the exact kEpsilon weight. + constexpr double kEpsilon = 0.001; + double numValues = state.valueLevelVars.size(); + double levelSpan = std::max(1, state.bootstrapWaterline); + double tieBreakCeiling = 0.5 * minLevelMarginalCost(costModel); + double perValueWeight = kEpsilon; + if (kEpsilon * numValues * levelSpan > tieBreakCeiling) { + perValueWeight = tieBreakCeiling / (numValues * levelSpan); + } for (auto& [value, levelVar] : state.valueLevelVars) { - objective += -0.001 * levelVar; + objective += -perValueWeight * levelVar; + } + // At a partition boundary, prefer low output scale as a proxy for the + // rescale work pushed onto downstream partitions (Orbit's 0.2-weighted + // boundary term). Excluded from the recorded partition cost. + if (state.boundary.outputValue && !state.levelOnly) { + objective += + 0.2 * costModel.rescaleCost * state.boundary.sizeInOps * + state.valueScaleVars.at(state.canon(state.boundary.outputValue)); } state.model.Minimize(objective); } -static void populateSolution( - const math_opt::SolveResult& result, ILPModelState& state, - llvm::DenseMap& solution, - llvm::DenseMap& solutionLevelBeforeBootstrap, - llvm::DenseMap& solutionLevelAfterBootstrap, - llvm::SmallVector& - nodeManagement, - llvm::SmallVector& - edgeManagement) { +// Extract the mgmt decisions for every member op of every group in range, and +// re-evaluate the exact objective share of this partition (excluding +// tie-break and boundary-pressure terms). +static void populatePartitionSolution(const math_opt::SolveResult& result, + ILPModelState& state, + const OpCostModel& costModel, + PartitionSolution& soln) { auto varMap = result.variable_values(); - for (Operation* op : state.trackedOps) { - bool useBootstrap = varMap.at(state.bootstrapVars.at(op)) > 0.5; - solution.insert(std::make_pair(op, useBootstrap)); - - int inputLevel = roundedValue(varMap, state.inputLevelVars.at(op)); - int inputScale = roundedValue(varMap, state.inputScaleVars.at(op)); - for (OpResult result : op->getResults()) { - if (!isSecret(result, state.solver)) continue; - int outputLevel = roundedValue(varMap, state.valueLevelVars.at(result)); - int outputScale = roundedValue(varMap, state.valueScaleVars.at(result)); - solutionLevelBeforeBootstrap.insert(std::make_pair(result, inputLevel)); - solutionLevelAfterBootstrap.insert(std::make_pair(result, outputLevel)); - nodeManagement.push_back({result, inputLevel, inputScale, outputLevel, - outputScale, useBootstrap}); + double cost = 0; + for (int gi = state.groupBegin; gi < state.groupEnd; ++gi) { + const OpGroup& group = state.grouping.groups[gi]; + Operation* rep = group.representative; + bool useBootstrap = varMap.at(state.bootstrapVars.at(rep)) > 0.5; + int nodeRescales = roundedValue(varMap, state.nodeRescaleVars.at(rep)); + cost += group.weight * costModel.bootstrapCost * (useBootstrap ? 1 : 0); + cost += group.weight * costModel.rescaleCost * nodeRescales; + + int inputLevel = roundedValue(varMap, state.inputLevelVars.at(rep)); + int inputScale = roundedValue(varMap, state.inputScaleVars.at(rep)); + llvm::DenseSet interior(group.interiorValues.begin(), + group.interiorValues.end()); + for (Operation* member : group.members) { + bool isManagementSite = false; + for (OpResult result : member->getResults()) { + if (!isSecret(result, state.solver)) continue; + if (interior.contains(result)) { + // Addition-tree interior values stay at the group's input state. + soln.nodeManagement.push_back( + {result, inputLevel, inputScale, inputLevel, inputScale, false}); + soln.levelBefore.insert(std::make_pair(result, inputLevel)); + soln.levelAfter.insert(std::make_pair(result, inputLevel)); + continue; + } + isManagementSite = true; + Value key = state.canon(result); + int outputLevel = roundedValue(varMap, state.valueLevelVars.at(key)); + int outputScale = roundedValue(varMap, state.valueScaleVars.at(key)); + soln.nodeManagement.push_back({result, inputLevel, inputScale, + outputLevel, outputScale, useBootstrap}); + soln.levelBefore.insert(std::make_pair(result, inputLevel)); + soln.levelAfter.insert(std::make_pair(result, outputLevel)); + } + soln.bootstrapDecisions.insert( + std::make_pair(member, isManagementSite && useBootstrap)); } } + for (const GroupEdge& edge : state.edges) { + cost += edge.weight * costModel.rescaleCost * + roundedValue(varMap, edge.rescaleVar); + } + if (costModel.hasLevelCosts) { + for (Operation* op : state.trackedOps) { + auto opCost = levelCostForOp(op, state.solver, costModel); + if (!opCost.has_value()) continue; + cost += + opCost->slope * roundedValue(varMap, state.inputLevelVars.at(op)) + + opCost->intercept; + } + } + soln.cost = cost; for (Operation* op : state.trackedOps) { int targetLevel = roundedValue(varMap, state.inputLevelVars.at(op)); + // One transition per (op, producer class): repeated operand slots consume + // the same managed value, so the decoder rescales the producer once and + // points every matching slot at it. + DenseSet emittedForOp; for (OpOperand& operandUse : op->getOpOperands()) { if (!state.edgeScaleVars.contains(&operandUse)) continue; - Value operand = operandUse.get(); - int sourceLevel = roundedValue(varMap, state.valueLevelVars.at(operand)); - int sourceScale = roundedValue(varMap, state.valueScaleVars.at(operand)); + Value key = state.canon(operandUse.get()); + if (!emittedForOp.insert(key).second) continue; + int sourceLevel = roundedValue(varMap, state.valueLevelVars.at(key)); + int sourceScale = roundedValue(varMap, state.valueScaleVars.at(key)); int targetScale = roundedValue(varMap, state.edgeScaleVars.at(&operandUse)); - edgeManagement.push_back({op, operandUse.getOperandNumber(), sourceLevel, - sourceScale, targetLevel, targetScale}); + soln.edgeManagement.push_back({op, operandUse.getOperandNumber(), + sourceLevel, sourceScale, targetLevel, + targetScale}); } } -} - -LogicalResult ILPBootstrapPlacementAnalysis::solve() { - auto genericOp = dyn_cast(opToRunOn); - if (!genericOp) return failure(); - ILPModelState state(genericOp.getBody(), solver, bootstrapWaterline, - scaleWaterline, scaleFactorBits, scaleMode); + if (state.boundary.outputValue) { + Value key = state.canon(state.boundary.outputValue); + soln.outState.level = roundedValue(varMap, state.valueLevelVars.at(key)); + soln.outState.scale = roundedValue(varMap, state.valueScaleVars.at(key)); + } +} - if (failed(addBodyArgumentVariables(state))) return failure(); - addTrackedOperationVariables(state); - addOperandEdgeConstraints(state); +// Build and solve the model for one partition under one boundary state. +// Returns failure only on hard errors; an infeasible boundary state leaves +// solutionOut empty. +static LogicalResult solvePartition( + Block* body, DataFlowSolver* solver, const Options& options, + const OpGrouping& grouping, int groupBegin, int groupEnd, + const PartitionBoundary& boundary, + std::optional& solutionOut) { + solutionOut.reset(); + ILPModelState state(body, solver, options, grouping, groupBegin, groupEnd, + boundary); + + addPinnedInputVariables(state); + addTrackedGroupVariables(state); + addGroupEdgeConstraints(state); if (failed(addYieldConstraints(state))) return failure(); - addNodeTransitionConstraints(state, bootstrapLevelLowerBound); - addObjective(state, costModel); + addNodeTransitionConstraints(state, options.bootstrapLevelLowerBound); + addOutputPinConstraints(state); + addObjective(state, options.costModel); LLVM_DEBUG({ std::stringstream ss; @@ -574,15 +774,297 @@ LogicalResult ILPBootstrapPlacementAnalysis::solve() { case math_opt::TerminationReason::kFeasible: break; default: + // Infeasible under this boundary state; the caller decides whether that + // is an error. + LLVM_DEBUG(llvm::dbgs() + << "partition solve infeasible, termination status code: " + << static_cast(result.termination.reason) << "\n"); + return success(); + } + + PartitionSolution soln; + populatePartitionSolution(result, state, options.costModel, soln); + solutionOut = std::move(soln); + return success(); +} + +// Cut the group sequence at single-input single-output boundaries: positions +// where exactly one value class is live across the cut (Orbit's SISO +// partitioning). +static SmallVector computePartitions(const OpGrouping& grouping, + Block* body, + DataFlowSolver* solver, + int partitionMinSize) { + int numGroups = grouping.groups.size(); + SmallVector partitions; + if (numGroups == 0) { + partitions.push_back({0, 0, Value(), 0}); + return partitions; + } + + // Producer position and last consumer position per canonical value class. + llvm::DenseMap producerPos; + llvm::DenseMap lastUsePos; + for (BlockArgument arg : body->getArguments()) { + if (isSecret(arg, solver)) producerPos[arg] = -1; + } + for (int gi = 0; gi < numGroups; ++gi) { + for (Value result : grouping.groups[gi].resultValues) { + producerPos[grouping.canonicalValue(result)] = gi; + } + } + for (int gi = 0; gi < numGroups; ++gi) { + for (Operation* member : grouping.groups[gi].members) { + for (Value operand : member->getOperands()) { + if (!isSecret(operand, solver)) continue; + if (Operation* def = operand.getDefiningOp()) { + auto it = grouping.groupIdOf.find(def); + if (it != grouping.groupIdOf.end() && it->second == gi) continue; + } + Value key = grouping.canonicalValue(operand); + if (!producerPos.contains(key)) continue; + int& last = lastUsePos[key]; + last = std::max(last, gi); + } + } + } + for (Operation& op : body->getOperations()) { + if (!isa(op)) continue; + for (Value operand : op.getOperands()) { + Value key = grouping.canonicalValue(operand); + if (producerPos.contains(key)) lastUsePos[key] = INT_MAX; + } + } + + // Live-class count at each candidate cut (between positions p and p + 1). + std::vector delta(numGroups + 1, 0); + for (auto& [key, prod] : producerPos) { + auto it = lastUsePos.find(key); + if (it == lastUsePos.end() || it->second <= prod) continue; + int lo = std::max(prod, 0); + int hi = std::min(it->second == INT_MAX ? numGroups - 1 : it->second - 1, + numGroups - 2); + if (lo > hi) continue; + delta[lo] += 1; + delta[hi + 1] -= 1; + } + auto cutValueAt = [&](int p) -> Value { + for (auto& [key, prod] : producerPos) { + auto it = lastUsePos.find(key); + if (it == lastUsePos.end()) continue; + if (prod <= p && it->second > p) return key; + } + return Value(); + }; + + SmallVector cuts; + int live = 0; + for (int p = 0; p <= numGroups - 2; ++p) { + live += delta[p]; + if (live == 1) cuts.push_back(p); + } + + // Enforce the minimum partition size in original ops (Orbit's delta). + std::vector prefixOps(numGroups + 1, 0); + for (int gi = 0; gi < numGroups; ++gi) { + prefixOps[gi + 1] = prefixOps[gi] + grouping.groups[gi].members.size(); + } + int minSize = std::max(1, partitionMinSize); + SmallVector kept; + int begin = 0; + for (int p : cuts) { + if (prefixOps[p + 1] - prefixOps[begin] >= minSize) { + kept.push_back(p); + begin = p + 1; + } + } + while (!kept.empty() && + prefixOps[numGroups] - prefixOps[kept.back() + 1] < minSize) { + kept.pop_back(); + } + + begin = 0; + for (int p : kept) { + partitions.push_back( + {begin, p + 1, cutValueAt(p), prefixOps[p + 1] - prefixOps[begin]}); + begin = p + 1; + } + partitions.push_back( + {begin, numGroups, Value(), prefixOps[numGroups] - prefixOps[begin]}); + return partitions; +} + +LogicalResult ILPBootstrapPlacementAnalysis::solve() { + auto genericOp = dyn_cast(opToRunOn); + if (!genericOp) return failure(); + Block* body = genericOp.getBody(); + + OpGrouping grouping = computeOpGrouping(body, solver, options.compress); + SmallVector, 4> argStates; + if (failed(computeArgInitStates(body, solver, options, argStates))) + return failure(); + SmallVector partitions = + computePartitions(grouping, body, solver, options.partitionMinSize); + + LLVM_DEBUG({ + int numOps = 0; + for (const OpGroup& group : grouping.groups) numOps += group.members.size(); + llvm::dbgs() << "ilp-bootstrap-placement: " << numOps << " tracked ops in " + << grouping.groups.size() << " groups across " + << partitions.size() << " partitions\n"; + }); + + auto adoptSolution = [&](PartitionSolution& soln) { + for (auto& management : soln.nodeManagement) + nodeManagement.push_back(management); + for (auto& management : soln.edgeManagement) + edgeManagement.push_back(management); + for (auto& [op, useBootstrap] : soln.bootstrapDecisions) + solution.insert(std::make_pair(op, useBootstrap)); + for (auto& [value, level] : soln.levelBefore) + solutionLevelBeforeBootstrap.insert(std::make_pair(value, level)); + for (auto& [value, level] : soln.levelAfter) + solutionLevelAfterBootstrap.insert(std::make_pair(value, level)); + }; + + if (partitions.size() == 1) { + PartitionBoundary boundary; + boundary.pinnedInputs = argStates; + boundary.applyYieldConstraints = true; + boundary.sizeInOps = partitions[0].sizeInOps; + std::optional soln; + if (failed(solvePartition(body, solver, options, grouping, 0, + grouping.groups.size(), boundary, soln))) + return failure(); + if (!soln.has_value()) { llvm::errs() << "No feasible solution found (the problem may be " - "infeasible). Termination status code: " - << static_cast(result.termination.reason) << "\n"; + "infeasible).\n"; return failure(); + } + adoptSolution(*soln); + return success(); } - populateSolution(result, state, solution, solutionLevelBeforeBootstrap, - solutionLevelAfterBootstrap, nodeManagement, edgeManagement); + // Dynamic program over partitions (Orbit's QBP + DP stitch): each partition + // is solved for every reachable boundary input state and every enumerated + // boundary output level; the realized output scale keys the next + // partition's input states. + using StateKey = std::pair; + struct DpEntry { + double totalCost; + StateKey prevState; + int solutionIndex; + }; + // Sentinel input for the first partition (arguments are pinned separately) + // and sentinel output for the last. + const StateKey kNoState(-1, -1); + std::map prevTable; + prevTable.insert({kNoState, {0.0, kNoState, -1}}); + std::vector> solutions(partitions.size()); + std::vector> tables(partitions.size()); + + for (size_t k = 0; k < partitions.size(); ++k) { + const Partition& partition = partitions[k]; + bool isLast = k == partitions.size() - 1; + std::map table; + + for (auto& [inState, prevEntry] : prevTable) { + PartitionBoundary boundary; + if (k == 0) { + boundary.pinnedInputs = argStates; + } else { + boundary.pinnedInputs.push_back( + {partitions[k - 1].cutValue, {inState.first, inState.second}}); + } + boundary.applyYieldConstraints = isLast; + boundary.sizeInOps = partition.sizeInOps; + + SmallVector outLevels; + if (isLast) { + outLevels.push_back(-1); + } else { + boundary.outputValue = partition.cutValue; + for (int level = 0; level <= options.bootstrapWaterline; ++level) + outLevels.push_back(level); + } + + for (int outLevel : outLevels) { + boundary.outputLevel = outLevel; + std::optional soln; + if (failed(solvePartition(body, solver, options, grouping, + partition.groupBegin, partition.groupEnd, + boundary, soln))) + return failure(); + LLVM_DEBUG(llvm::dbgs() + << "partition " << k << " in(" << inState.first << "," + << inState.second << ") outLvl " << outLevel << ": " + << (soln.has_value() + ? ("cost " + std::to_string(soln->cost) + " out(" + + std::to_string(soln->outState.level) + "," + + std::to_string(soln->outState.scale) + ")") + : std::string("infeasible")) + << "\n"); + if (!soln.has_value()) continue; + + StateKey outKey = + isLast ? kNoState + : StateKey(soln->outState.level, soln->outState.scale); + double totalCost = prevEntry.totalCost + soln->cost; + auto it = table.find(outKey); + if (it != table.end() && it->second.totalCost <= totalCost) continue; + soln->inKey = inState; + int solutionIndex = solutions[k].size(); + solutions[k].push_back(std::move(*soln)); + table[outKey] = {totalCost, inState, solutionIndex}; + } + } + + if (table.empty()) { + genericOp->emitError() + << "no feasible boundary state found for ILP partition " << k; + return failure(); + } + + // Prune to at most one boundary scale per boundary level: keep the + // minimal scale whose cost is within tolerance of the level's best. + if (!isLast) { + double tolerance = + 0.8 * partition.sizeInOps * options.costModel.rescaleCost; + std::map bestCostPerLevel; + for (auto& [key, entry] : table) { + auto it = bestCostPerLevel.find(key.first); + if (it == bestCostPerLevel.end() || entry.totalCost < it->second) + bestCostPerLevel[key.first] = entry.totalCost; + } + std::map pruned; + std::map doneLevel; + for (auto& [key, entry] : table) { + // std::map iterates scales in increasing order per level, so the + // first qualifying scale for a level wins. + if (doneLevel[key.first]) continue; + if (entry.totalCost <= bestCostPerLevel[key.first] + tolerance) { + pruned.insert({key, entry}); + doneLevel[key.first] = true; + } + } + table = std::move(pruned); + } + tables[k] = table; + prevTable = std::move(table); + } + + // Backtrack from the final sentinel state and stitch the chosen solutions. + StateKey key = kNoState; + SmallVector chosen(partitions.size()); + for (int k = partitions.size() - 1; k >= 0; --k) { + const DpEntry& entry = tables[k].at(key); + chosen[k] = entry.solutionIndex; + key = entry.prevState; + } + for (size_t k = 0; k < partitions.size(); ++k) { + adoptSolution(solutions[k][chosen[k]]); + } return success(); } @@ -596,9 +1078,9 @@ void ILPBootstrapPlacementAnalysis::printSolution(llvm::raw_ostream& os) const { if (!body) return; os << "--- ILP bootstrap placement solution ---\n"; - os << "bootstrap waterline: " << bootstrapWaterline << "\n"; - os << "scale waterline: " << scaleWaterline << "\n"; - os << "scale factor bits: " << scaleFactorBits << "\n\n"; + os << "bootstrap waterline: " << options.bootstrapWaterline << "\n"; + os << "scale waterline: " << options.scaleWaterline << "\n"; + os << "scale factor bits: " << options.scaleFactorBits << "\n\n"; for (BlockArgument arg : body->getArguments()) { auto it = solutionLevelAfterBootstrap.find(arg); diff --git a/lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.h b/lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.h index ac5918a0e3..1b86a600a4 100644 --- a/lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.h +++ b/lib/Analysis/ILPBootstrapPlacementAnalysis/ILPBootstrapPlacementAnalysis.h @@ -43,6 +43,23 @@ 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; @@ -62,19 +79,8 @@ class ILPBootstrapPlacementAnalysis { }; ILPBootstrapPlacementAnalysis(Operation* op, DataFlowSolver* solver, - int bootstrapWaterline, int scaleWaterline, - int scaleFactorBits, - int bootstrapLevelLowerBound, - const OpCostModel& costModel, - ScaleMode scaleMode) - : opToRunOn(op), - solver(solver), - bootstrapWaterline(bootstrapWaterline), - scaleWaterline(scaleWaterline), - scaleFactorBits(scaleFactorBits), - bootstrapLevelLowerBound(bootstrapLevelLowerBound), - costModel(costModel), - scaleMode(scaleMode) {} + const Options& options) + : opToRunOn(op), solver(solver), options(options) {} ~ILPBootstrapPlacementAnalysis() = default; LogicalResult solve(); @@ -107,12 +113,7 @@ class ILPBootstrapPlacementAnalysis { private: Operation* opToRunOn; DataFlowSolver* solver; - int bootstrapWaterline; - int scaleWaterline; - int scaleFactorBits; - int bootstrapLevelLowerBound; - OpCostModel costModel; - ScaleMode scaleMode; + Options options; llvm::DenseMap solution; llvm::DenseMap solutionLevelBeforeBootstrap; llvm::DenseMap solutionLevelAfterBootstrap; diff --git a/lib/Analysis/ILPBootstrapPlacementAnalysis/OpGrouping.cpp b/lib/Analysis/ILPBootstrapPlacementAnalysis/OpGrouping.cpp new file mode 100644 index 0000000000..53600cee8d --- /dev/null +++ b/lib/Analysis/ILPBootstrapPlacementAnalysis/OpGrouping.cpp @@ -0,0 +1,344 @@ +#include "lib/Analysis/ILPBootstrapPlacementAnalysis/OpGrouping.h" + +#include +#include +#include +#include +#include + +#include "lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.h" +#include "lib/Dialect/Secret/IR/SecretOps.h" +#include "llvm/include/llvm/ADT/DenseMap.h" // from @llvm-project +#include "llvm/include/llvm/ADT/DenseSet.h" // from @llvm-project +#include "llvm/include/llvm/ADT/STLExtras.h" // from @llvm-project +#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project + +namespace mlir { +namespace heir { + +bool isMultiplication(Operation* op) { + return isa(op); +} + +bool isAdditionLike(Operation* op) { + return isa(op); +} + +bool isConstantLike(Value value) { + return value.getDefiningOp() != nullptr; +} + +static bool hasSecretResult(Operation& op, DataFlowSolver* solver) { + return llvm::any_of(op.getResults(), [&](OpResult result) { + return isSecret(result, solver); + }); +} + +bool shouldTrackOperation(Operation& op, DataFlowSolver* solver) { + return !isa(op) && hasSecretResult(op, solver); +} + +namespace { + +struct GroupingContext { + GroupingContext(Block* body, DataFlowSolver* solver) + : body(body), solver(solver) {} + + Block* body; + DataFlowSolver* solver; + SmallVector trackedOps; + // Program-order index and longest-path depth of each tracked op. + DenseMap opIndex; + DenseMap opDepth; +}; + +} // namespace + +static OpGroup makeSingletonGroup(GroupingContext& ctx, Operation* op) { + OpGroup group; + group.members.push_back(op); + group.representative = op; + for (OpResult result : op->getResults()) { + if (isSecret(result, ctx.solver)) group.resultValues.push_back(result); + } + group.weight = 1; + group.isMultiplication = isMultiplication(op); + group.depth = ctx.opDepth.lookup(op); + return group; +} + +// Collapse maximal addition trees into one group each (Orbit's +// addition_squash). All additions in a tree must execute at the same level +// and scale, so they need only one set of ILP variables and one management +// decision after the final (deepest) addition. An addition joins the tree of +// a consumer only when its entire fanout is inside that tree, so no interior +// sum escapes at a possibly different state. +static SmallVector squashAdditionTrees(GroupingContext& ctx, + bool compress) { + auto isSquashable = [&](Operation* op) { + return compress && isAdditionLike(op) && op->getNumResults() == 1; + }; + + // Grow trees from their deepest member so producers join consumers. + SmallVector order(ctx.trackedOps); + llvm::stable_sort(order, [&](Operation* a, Operation* b) { + return ctx.opDepth.lookup(a) > ctx.opDepth.lookup(b); + }); + + SmallVector groups; + DenseSet used; + for (Operation* head : order) { + if (used.contains(head)) continue; + if (!isSquashable(head)) { + used.insert(head); + groups.push_back(makeSingletonGroup(ctx, head)); + continue; + } + + // Fixpoint: absorb an addition when all users of its result are already + // in the tree. Absorbing one addition can make its producers eligible. + DenseSet tree; + tree.insert(head); + bool changed = true; + while (changed) { + changed = false; + SmallVector candidates; + for (Operation* member : tree) { + for (Value operand : member->getOperands()) { + Operation* def = operand.getDefiningOp(); + if (!def || def->getBlock() != ctx.body) continue; + if (used.contains(def) || tree.contains(def)) continue; + if (!ctx.opIndex.contains(def) || !isSquashable(def)) continue; + candidates.push_back(def); + } + } + for (Operation* candidate : candidates) { + if (tree.contains(candidate)) continue; + bool hasUser = false; + bool allUsersInside = true; + for (Operation* user : candidate->getResult(0).getUsers()) { + hasUser = true; + if (!tree.contains(user)) { + allUsersInside = false; + break; + } + } + if (hasUser && allUsersInside) { + tree.insert(candidate); + changed = true; + } + } + } + + OpGroup group; + group.members.append(tree.begin(), tree.end()); + llvm::stable_sort(group.members, [&](Operation* a, Operation* b) { + return ctx.opIndex.lookup(a) < ctx.opIndex.lookup(b); + }); + group.representative = head; + group.resultValues.push_back(head->getResult(0)); + for (Operation* member : group.members) { + used.insert(member); + if (member != head) group.interiorValues.push_back(member->getResult(0)); + } + group.weight = 1; + group.isMultiplication = false; + group.depth = ctx.opDepth.lookup(head); + groups.push_back(std::move(group)); + } + return groups; +} + +// Merge structurally equivalent groups by compression (Orbit's auto_compress): +// iterative label refinement, where each round relabels a group by its op +// class, depth, and the current labels of its producers and consumers until +// the labeling reaches a fixpoint. Two groups may merge only when they have +// the same depth, the same op class, and — at the fixpoint — the same set of +// producer classes; that guarantees replaying one representative's solution on +// every member satisfies each member's constraints, and cost linearity makes +// the merged objective exactly the sum over members. +static void mergeEquivalentGroups(GroupingContext& ctx, + SmallVector& groups) { + int numGroups = groups.size(); + if (numGroups <= 1) return; + + DenseMap groupOf; + for (auto [i, group] : llvm::enumerate(groups)) { + for (Operation* member : group.members) groupOf[member] = i; + } + + // Groups whose results are yielded (or that have multiple results) never + // merge: yielded values can carry per-result mgmt.mgmt pins that must not + // propagate to other ops. + std::vector uniqueTag(numGroups, 0); + int nextUniqueTag = 1; + for (int i = 0; i < numGroups; ++i) { + bool unique = groups[i].representative->getNumResults() != 1; + for (Value result : groups[i].resultValues) { + for (Operation* user : result.getUsers()) { + if (isa(user)) unique = true; + } + } + if (unique) uniqueTag[i] = nextUniqueTag++; + } + + // Operand labels are (kind, id) pairs so distinct kinds cannot collide: + // kind 0 = producer group (id = current-round label), 1 = secret block + // argument, 2 = constant, 3 = other non-secret value, 4 = secret value not + // produced by a tracked op (never merged across). + using OperandLabel = std::pair; + DenseMap externalIds; + + int addOpKey = 0; + std::map opKeys; + auto opKeyOf = [&](const OpGroup& group) { + if (isAdditionLike(group.representative)) return addOpKey; + auto [it, inserted] = opKeys.try_emplace( + group.representative->getName().getStringRef(), opKeys.size() + 1); + return it->second; + }; + + std::vector order(numGroups); + for (int i = 0; i < numGroups; ++i) order[i] = i; + llvm::stable_sort(order, [&](int a, int b) { + if (groups[a].depth != groups[b].depth) + return groups[a].depth < groups[b].depth; + return ctx.opIndex.lookup(groups[a].members.front()) < + ctx.opIndex.lookup(groups[b].members.front()); + }); + + using Descriptor = + std::tuple, std::vector>; + std::vector labels(numGroups, 0); + int numDistinct = 1; + constexpr int kMaxRounds = 100; + for (int round = 0; round < kMaxRounds; ++round) { + std::vector newLabels(numGroups, 0); + std::map intern; + for (int i : order) { + const OpGroup& group = groups[i]; + + // Producer classes, from this round's labels (producers are strictly + // shallower, so they are relabeled before their consumers). + std::vector inputs; + for (Operation* member : group.members) { + for (Value operand : member->getOperands()) { + Operation* def = operand.getDefiningOp(); + if (def && groupOf.count(def)) { + if (groupOf[def] == i) continue; // interior edge + inputs.emplace_back(0, newLabels[groupOf[def]]); + continue; + } + if (auto arg = dyn_cast(operand); + arg && arg.getOwner() == ctx.body && + isSecret(operand, ctx.solver)) { + inputs.emplace_back(1, arg.getArgNumber()); + continue; + } + if (isConstantLike(operand)) { + inputs.emplace_back(2, 0); + continue; + } + if (!isSecret(operand, ctx.solver)) { + inputs.emplace_back(3, 0); + continue; + } + auto [it, inserted] = + externalIds.try_emplace(operand, externalIds.size()); + inputs.emplace_back(4, it->second); + } + } + llvm::sort(inputs); + inputs.erase(std::unique(inputs.begin(), inputs.end()), inputs.end()); + + // Consumer classes, from the previous round's labels. + std::vector successors; + for (Value result : group.resultValues) { + for (Operation* user : result.getUsers()) { + if (groupOf.count(user)) successors.push_back(labels[groupOf[user]]); + } + } + llvm::sort(successors); + successors.erase(std::unique(successors.begin(), successors.end()), + successors.end()); + + Descriptor descriptor(group.depth, opKeyOf(group), uniqueTag[i], + std::move(inputs), std::move(successors)); + auto [it, inserted] = + intern.try_emplace(std::move(descriptor), intern.size()); + newLabels[i] = it->second; + } + int newDistinct = intern.size(); + labels = std::move(newLabels); + if (newDistinct == numDistinct) break; + numDistinct = newDistinct; + } + + if (numDistinct == numGroups) return; // nothing merged + + std::map> buckets; + for (int i = 0; i < numGroups; ++i) buckets[labels[i]].push_back(i); + + SmallVector merged; + for (auto& [label, bucket] : buckets) { + OpGroup group = std::move(groups[bucket.front()]); + for (size_t k = 1; k < bucket.size(); ++k) { + OpGroup& other = groups[bucket[k]]; + group.members.append(other.members.begin(), other.members.end()); + group.resultValues.append(other.resultValues.begin(), + other.resultValues.end()); + group.interiorValues.append(other.interiorValues.begin(), + other.interiorValues.end()); + group.weight += other.weight; + } + llvm::stable_sort(group.members, [&](Operation* a, Operation* b) { + return ctx.opIndex.lookup(a) < ctx.opIndex.lookup(b); + }); + merged.push_back(std::move(group)); + } + groups = std::move(merged); +} + +OpGrouping computeOpGrouping(Block* body, DataFlowSolver* solver, + bool compress) { + GroupingContext ctx(body, solver); + int index = 0; + for (Operation& op : body->getOperations()) { + if (!shouldTrackOperation(op, solver)) continue; + ctx.trackedOps.push_back(&op); + ctx.opIndex[&op] = index++; + int depth = 0; + for (Value operand : op.getOperands()) { + Operation* def = operand.getDefiningOp(); + auto it = ctx.opDepth.find(def); + if (def && it != ctx.opDepth.end()) + depth = std::max(depth, it->second + 1); + } + ctx.opDepth[&op] = depth; + } + + SmallVector groups = squashAdditionTrees(ctx, compress); + if (compress) mergeEquivalentGroups(ctx, groups); + + llvm::stable_sort(groups, [&](const OpGroup& a, const OpGroup& b) { + if (a.depth != b.depth) return a.depth < b.depth; + return ctx.opIndex.lookup(a.members.front()) < + ctx.opIndex.lookup(b.members.front()); + }); + + OpGrouping grouping; + grouping.groups = std::move(groups); + for (auto [i, group] : llvm::enumerate(grouping.groups)) { + for (Operation* member : group.members) grouping.groupIdOf[member] = i; + if (group.representative->getNumResults() == 1) { + Value repResult = group.representative->getResult(0); + for (Value result : group.resultValues) { + if (result != repResult) grouping.valueRep[result] = repResult; + } + } + } + return grouping; +} + +} // namespace heir +} // namespace mlir diff --git a/lib/Analysis/ILPBootstrapPlacementAnalysis/OpGrouping.h b/lib/Analysis/ILPBootstrapPlacementAnalysis/OpGrouping.h new file mode 100644 index 0000000000..58984a187e --- /dev/null +++ b/lib/Analysis/ILPBootstrapPlacementAnalysis/OpGrouping.h @@ -0,0 +1,69 @@ +#ifndef LIB_ANALYSIS_ILPBOOTSTRAPPLACEMENTANALYSIS_OPGROUPING_H_ +#define LIB_ANALYSIS_ILPBOOTSTRAPPLACEMENTANALYSIS_OPGROUPING_H_ + +#include "llvm/include/llvm/ADT/DenseMap.h" // from @llvm-project +#include "llvm/include/llvm/ADT/SmallVector.h" // from @llvm-project +#include "mlir/include/mlir/Analysis/DataFlowFramework.h" // from @llvm-project +#include "mlir/include/mlir/IR/Block.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/Support/LLVM.h" // from @llvm-project + +namespace mlir { +namespace heir { + +bool isMultiplication(Operation* op); +bool isAdditionLike(Operation* op); +bool isConstantLike(Value value); +bool shouldTrackOperation(Operation& op, DataFlowSolver* solver); + +// One ILP decision class: a set of ops that share one set of decision +// variables (input level/scale, rescale count, bootstrap) and one output +// state. Grouping ops shrinks the ILP without changing its optimum. +struct OpGroup { + // All member ops, in program order. + SmallVector members; + // The member whose operands and results define the group's constraint + // structure and variable names. + Operation* representative = nullptr; + // Values that carry the group's output state: the result of each merged + // management site. Management chosen for the group is decoded after each of + // these values. + SmallVector resultValues; + // Addition-tree interior results. These are consumed only inside the group + // and stay at the group's input state; no management is decoded for them. + SmallVector interiorValues; + // Number of management sites merged into this group: bootstrap and rescale + // decisions are charged weight times in the objective and decoded after + // each value in resultValues. + int weight = 1; + bool isMultiplication = false; + // Longest-path depth of the deepest member; group order in OpGrouping is + // topological by this depth. + int depth = 0; +}; + +struct OpGrouping { + // Topologically ordered by (depth, program order). + SmallVector groups; + // Tracked op -> index into groups. + DenseMap groupIdOf; + // Maps a group-output value to the representative value whose ILP variables + // it shares. Values absent from the map are their own representative. + DenseMap valueRep; + + Value canonicalValue(Value value) const { + auto it = valueRep.find(value); + return it == valueRep.end() ? value : it->second; + } +}; + +// Group the tracked operations of a secret.generic body. With compress false, +// every op is its own group. +OpGrouping computeOpGrouping(Block* body, DataFlowSolver* solver, + bool compress); + +} // namespace heir +} // namespace mlir + +#endif // LIB_ANALYSIS_ILPBOOTSTRAPPLACEMENTANALYSIS_OPGROUPING_H_ diff --git a/lib/Transforms/ILPBootstrapPlacement/ILPBootstrapPlacement.cpp b/lib/Transforms/ILPBootstrapPlacement/ILPBootstrapPlacement.cpp index a3cce02936..7fa0f003d0 100644 --- a/lib/Transforms/ILPBootstrapPlacement/ILPBootstrapPlacement.cpp +++ b/lib/Transforms/ILPBootstrapPlacement/ILPBootstrapPlacement.cpp @@ -227,10 +227,16 @@ struct ILPBootstrapPlacement effectiveCostModel = *loadedCostModel; } - ILPBootstrapPlacementAnalysis analysis( - genericOp, solver, bootstrapWaterline, scaleConfig.scaleWaterline, - scaleConfig.scaleFactorBits, bootstrapLevelLowerBound, - effectiveCostModel, scaleConfig.analysisScaleMode()); + ILPBootstrapPlacementAnalysis::Options analysisOptions; + analysisOptions.bootstrapWaterline = bootstrapWaterline; + analysisOptions.scaleWaterline = scaleConfig.scaleWaterline; + analysisOptions.scaleFactorBits = scaleConfig.scaleFactorBits; + analysisOptions.bootstrapLevelLowerBound = bootstrapLevelLowerBound; + analysisOptions.compress = compress; + analysisOptions.partitionMinSize = partitionMinSize; + analysisOptions.costModel = effectiveCostModel; + analysisOptions.scaleMode = scaleConfig.analysisScaleMode(); + ILPBootstrapPlacementAnalysis analysis(genericOp, solver, analysisOptions); if (failed(analysis.solve())) { genericOp->emitError( "Failed to solve the bootstrap placement optimization problem"); @@ -463,7 +469,9 @@ struct ILPBootstrapPlacement b, op, current, placement.inputLevel, placement.inputScale, placement.outputLevel, placement.outputScale, scaleConfig); if (failed(managed)) return failure(); - op->setOperand(placement.operandNumber, *managed); + // Point every operand slot consuming this producer at the single managed + // value, so a repeated operand (e.g. squaring x*x) is rescaled once. + op->replaceUsesOfWith(current, *managed); return success(); } diff --git a/lib/Transforms/ILPBootstrapPlacement/ILPBootstrapPlacement.td b/lib/Transforms/ILPBootstrapPlacement/ILPBootstrapPlacement.td index da3f327353..5d10600740 100644 --- a/lib/Transforms/ILPBootstrapPlacement/ILPBootstrapPlacement.td +++ b/lib/Transforms/ILPBootstrapPlacement/ILPBootstrapPlacement.td @@ -53,6 +53,16 @@ def ILPBootstrapPlacement : Pass <"ilp-bootstrap-placement"> { "int", /*default=*/"0", "Minimum input level at which bootstrap is allowed in the Orbit-inspired scale constraints.">, + Option<"compress", + "compress", + "bool", + /*default=*/"true", + "Group structurally equivalent ops (Orbit's compression) so they share ILP variables. Grouped ops execute at one common level and share one management decision per merged site.">, + Option<"partitionMinSize", + "partition-min-size", + "int", + /*default=*/"100", + "Minimum number of ops per single-input single-output partition (Orbit's delta). Partitions are solved independently under enumerated boundary states and stitched by dynamic programming.">, Option<"orbitCostModel", "orbit-cost-model", "std::string", diff --git a/lib/Transforms/ILPBootstrapPlacement/README.md b/lib/Transforms/ILPBootstrapPlacement/README.md index 4e18150bbc..637f3ef4dd 100644 --- a/lib/Transforms/ILPBootstrapPlacement/README.md +++ b/lib/Transforms/ILPBootstrapPlacement/README.md @@ -21,6 +21,33 @@ any is missing: objective. `CtCt` is the ciphertext-ciphertext variant of a binary op, `CtPt` the ciphertext-plaintext variant. +## Model reduction + +Two Orbit techniques shrink the ILP before it is solved; both are exact with +respect to the objective because op costs are linear in the execution level. + +**Compression** (`compress`, default true) groups ops so each group shares one +set of ILP variables and one management decision per merged site: + +- *Addition squashing*: a maximal tree of additions whose interior fanout stays + inside the tree executes at a single (level, scale) and needs one management + decision after its final addition. Note: Addition-tree squashing applies only + to additions. Rotate-and-sum trees are deliberately not squashed; equivalent + rotations are reduced by structural merging instead. +- *Structural merging* (`auto_compress`): via iterative label refinement, ops + with the same depth, op class, and (at the fixpoint) the same producer classes + share one variable set. Each merged op still decodes its own management ops, + and the objective charges the group once per member, so the compressed model's + optimum equals the original optimum restricted to symmetric solutions. + +**SISO partitioning** (`partition-min-size`, default 100) cuts the circuit where +exactly one value is live across the cut. Each partition is solved independently +for every reachable boundary input state and every enumerated boundary output +level, and a dynamic program stitches the per-partition solutions. At most one +scale per boundary level survives between partitions, so the stitched placement +is a high-quality heuristic rather than provably optimal — matching Orbit's +implementation. + ## Solver configuration The ILP is solved to a fixed 1% relative optimality gap with no time limit, diff --git a/tests/Transforms/ilp_bootstrap_placement/orbit_compression.mlir b/tests/Transforms/ilp_bootstrap_placement/orbit_compression.mlir new file mode 100644 index 0000000000..bf71517ad9 --- /dev/null +++ b/tests/Transforms/ilp_bootstrap_placement/orbit_compression.mlir @@ -0,0 +1,35 @@ +// RUN: heir-opt --ilp-bootstrap-placement=bootstrap-waterline=2 %s | FileCheck %s +// RUN: heir-opt --ilp-bootstrap-placement="bootstrap-waterline=2 compress=false" %s | FileCheck %s + +// Compression must not change the chosen placement on this circuit: the +// three parallel muls of (%s, %input1) are structurally identical and merge +// into one compression class, and the addition chain squashes into one group. With +// bootstrap-waterline=2 the chain runs out of levels and exactly one +// bootstrap is optimal; a bootstrap on the merged three-mul class would be +// charged (and decoded) three times, so the solver avoids it. Both the +// compressed and uncompressed models produce identical IR. + +// CHECK: func.func @compression_invariance +// CHECK-COUNT-1: mgmt.bootstrap +// CHECK-NOT: mgmt.bootstrap + +!pt_ty = tensor<8xf32> +!ct_ty = !secret.secret + +func.func @compression_invariance( + %arg0: !ct_ty, %arg1: !ct_ty) -> !ct_ty { + %0 = secret.generic(%arg0: !ct_ty, %arg1: !ct_ty) { + ^body(%input0: !pt_ty, %input1: !pt_ty): + %s = arith.mulf %input0, %input0 : !pt_ty + // Three structurally identical ops: one compression class of weight 3. + %p1 = arith.mulf %s, %input1 : !pt_ty + %p2 = arith.mulf %s, %input1 : !pt_ty + %p3 = arith.mulf %s, %input1 : !pt_ty + // One addition tree: one squashed group. + %a1 = arith.addf %p1, %p2 : !pt_ty + %a2 = arith.addf %a1, %p3 : !pt_ty + %out = arith.mulf %a2, %a2 : !pt_ty + secret.yield %out : !pt_ty + } -> !ct_ty + return %0 : !ct_ty +} diff --git a/tests/Transforms/ilp_bootstrap_placement/orbit_partition_ckks.mlir b/tests/Transforms/ilp_bootstrap_placement/orbit_partition_ckks.mlir new file mode 100644 index 0000000000..179bdabcff --- /dev/null +++ b/tests/Transforms/ilp_bootstrap_placement/orbit_partition_ckks.mlir @@ -0,0 +1,33 @@ +// RUN: heir-opt --ilp-bootstrap-placement="bootstrap-waterline=3 scale-waterline=51 scale-factor-bits=51 orbit-cost-model=%S/orbit_cost_model.json partition-min-size=1" %s | FileCheck %s + +// The CKKS variant of partition stitching: boundary states are (level, scale) +// pairs, with the boundary scale realized by the solver under the +// output-scale pressure term rather than enumerated, so the stitched rescale +// positions may differ from the single-partition optimum on circuits this +// small (Orbit has the same property). What must hold: the annotated output +// state is reached with exactly two rescales and no bootstrap. + +// CHECK: func.func @partitioned_level_dependent_rescale +// CHECK-COUNT-2: mgmt.modreduce +// CHECK-NOT: mgmt.modreduce +// CHECK-NOT: mgmt.bootstrap + +!pt_ty = tensor<8xf32> +!ct_ty = !secret.secret + +module attributes {scheme.ckks} { + func.func @partitioned_level_dependent_rescale( + %arg0: !ct_ty) + -> (!ct_ty {mgmt.mgmt = #mgmt.mgmt}) { + %0 = secret.generic(%arg0: !ct_ty) { + ^body(%input0: tensor<8xf32>): + %m = arith.mulf %input0, %input0 : !pt_ty + %a1 = arith.addf %m, %m : !pt_ty + %a2 = arith.addf %a1, %a1 : !pt_ty + %a3 = arith.addf %a2, %a2 : !pt_ty + %mm = arith.mulf %a3, %a3 : !pt_ty + secret.yield %mm : !pt_ty + } -> (!ct_ty {mgmt.mgmt = #mgmt.mgmt}) + return %0 : !ct_ty + } +} diff --git a/tests/Transforms/ilp_bootstrap_placement/orbit_siso_partition.mlir b/tests/Transforms/ilp_bootstrap_placement/orbit_siso_partition.mlir new file mode 100644 index 0000000000..54946ca605 --- /dev/null +++ b/tests/Transforms/ilp_bootstrap_placement/orbit_siso_partition.mlir @@ -0,0 +1,29 @@ +// RUN: heir-opt --ilp-bootstrap-placement=bootstrap-waterline=3 %s | FileCheck %s +// RUN: heir-opt --ilp-bootstrap-placement="bootstrap-waterline=3 partition-min-size=1" %s | FileCheck %s +// RUN: heir-opt --ilp-bootstrap-placement="bootstrap-waterline=3 partition-min-size=1 compress=false" %s | FileCheck %s + +// A pure squaring chain has a single-input single-output cut after every +// mul, so partition-min-size=1 solves each op as its own partition and +// stitches the boundary levels by dynamic programming. The stitched solution +// must match the single-partition optimum: five squarings from level 3 need +// exactly one bootstrap. + +// CHECK: func.func @siso_partition_chain +// CHECK-COUNT-1: mgmt.bootstrap +// CHECK-NOT: mgmt.bootstrap + +!pt_ty = tensor<8xf32> +!ct_ty = !secret.secret + +func.func @siso_partition_chain(%arg0: !ct_ty) -> !ct_ty { + %0 = secret.generic(%arg0: !ct_ty) { + ^body(%input0: !pt_ty): + %m1 = arith.mulf %input0, %input0 : !pt_ty + %m2 = arith.mulf %m1, %m1 : !pt_ty + %m3 = arith.mulf %m2, %m2 : !pt_ty + %m4 = arith.mulf %m3, %m3 : !pt_ty + %m5 = arith.mulf %m4, %m4 : !pt_ty + secret.yield %m5 : !pt_ty + } -> !ct_ty + return %0 : !ct_ty +}