Support level reduction interface for multiple secret operands

lib/Analysis/LevelAnalysis/BootstrapWaterlineAnalysis.cpp

@@ -35,7 +35,8 @@ LogicalResult BootstrapWaterlineAnalysis::visitOperation(
   }
 
   // 2. Compute prospective level
-  LevelState prospectiveLevel = deriveResultLevel(op, operandLevelStates);
+  LevelState prospectiveLevel =
+      deriveResultLevel(op, operandLevelStates, &solverRef);
   if (levelBudget > 0 && prospectiveLevel.isInt() &&
       prospectiveLevel.getInt() > levelBudget) {
     LLVM_DEBUG({

lib/Analysis/LevelAnalysis/BootstrapWaterlineAnalysis.h

@@ -66,6 +66,7 @@ class BootstrapWaterlineAnalysis
                              int bootstrapLevelsConsumed = 0)
       : dataflow::SparseForwardDataFlowAnalysis<BootstrapWaterlineLattice>(
             solver),
+        solverRef(solver),
         waterline(waterline),
         levelBudget(levelBudget),
         bootstrapLevelsConsumed(bootstrapLevelsConsumed) {}
@@ -90,6 +91,7 @@ class BootstrapWaterlineAnalysis
   }
 
  private:
+  DataFlowSolver& solverRef;
   int waterline;
   int levelBudget;
   int bootstrapLevelsConsumed;

lib/Analysis/LevelAnalysis/LevelAnalysis.cpp

@@ -5,6 +5,7 @@
 #include <cstdint>
 #include <functional>
 #include <optional>
+#include <variant>
 
 #include "lib/Analysis/Utils.h"
 #include "lib/Dialect/HEIRInterfaces.h"
@@ -15,6 +16,7 @@
 #include "lib/Target/CompilationTarget/CompilationTarget.h"
 #include "lib/Utils/AttributeUtils.h"
 #include "lib/Utils/Utils.h"
+#include "llvm/include/llvm/ADT/STLExtras.h"               // from @llvm-project
 #include "llvm/include/llvm/ADT/TypeSwitch.h"              // from @llvm-project
 #include "llvm/include/llvm/Support/Debug.h"               // from @llvm-project
 #include "mlir/include/mlir/Analysis/DataFlowFramework.h"  // from @llvm-project
@@ -23,6 +25,7 @@
 #include "mlir/include/mlir/IR/BuiltinAttributes.h"        // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinOps.h"               // from @llvm-project
 #include "mlir/include/mlir/IR/BuiltinTypes.h"             // from @llvm-project
+#include "mlir/include/mlir/IR/Diagnostics.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/IR/Visitors.h"                 // from @llvm-project
@@ -53,18 +56,24 @@ namespace heir {
 };
 
 LevelState transferForward(ReducesLevelOpInterface op,
-                           ArrayRef<LevelState> operands) {
-  unsigned operandIdx = op.getOperandToReduce().getOperandNumber();
-  LevelState result = std::visit(
-      Overloaded{
-          [](MaxLevel) -> LevelState { return LevelState(Invalid{}); },
-          [](Uninit) -> LevelState { return LevelState(Uninit{}); },
-          [](Invalid) -> LevelState { return LevelState(Invalid{}); },
-          [&](int val) -> LevelState {
-            return LevelState(val + op.getLevelsToDrop());
-          },
-      },
-      operands[operandIdx].get());
+                           ArrayRef<LevelState> operands,
+                           const DataFlowSolver* solver) {
+  auto operandsToReduce = op.getOperandsToReduce(solver);
+  LevelState result;
+  for (auto* operand : operandsToReduce) {
+    unsigned operandIdx = operand->getOperandNumber();
+    LevelState opResult = std::visit(
+        Overloaded{
+            [](MaxLevel) -> LevelState { return LevelState(Invalid{}); },
+            [](Uninit) -> LevelState { return LevelState(Uninit{}); },
+            [](Invalid) -> LevelState { return LevelState(Invalid{}); },
+            [&](int val) -> LevelState {
+              return LevelState(val + op.getLevelsToDrop());
+            },
+        },
+        operands[operandIdx].get());
+    result = LevelState::join(result, opResult);
+  }
   LLVM_DEBUG(debugLog("ReduceLevelOpInterface", operands, result));
   return result;
 }
@@ -110,14 +119,16 @@ LevelState transferForward(ResetsLevelOpInterface op,
   return result;
 }
 
-LevelState deriveResultLevel(Operation* op, ArrayRef<LevelState> operands) {
+LevelState deriveResultLevel(Operation* op, ArrayRef<LevelState> operands,
+                             const DataFlowSolver* solver) {
   return llvm::TypeSwitch<Operation*, LevelState>(op)
       .Case<ResetsLevelOpInterface>(
           [&](auto op) -> LevelState { return transferForward(op, operands); })
       .Case<ReducesAllLevelsOpInterface>(
           [&](auto op) -> LevelState { return transferForward(op, operands); })
-      .Case<ReducesLevelOpInterface>(
-          [&](auto op) -> LevelState { return transferForward(op, operands); })
+      .Case<ReducesLevelOpInterface>([&](auto op) -> LevelState {
+        return transferForward(op, operands, solver);
+      })
       .Default([&](auto* op) -> LevelState {
         LevelState result;
         for (const auto& operand : operands) {
@@ -145,10 +156,32 @@ LogicalResult LevelAnalysis::visitOperation(
   for (auto* operand : operands) {
     operandStates.push_back(operand->getValue());
   }
-  LevelState resultLevel = deriveResultLevel(op, operandStates);
+  bool operandsValid = llvm::all_of(operandStates, [](const LevelState& state) {
+    return !state.isInvalid();
+  });
+  LevelState resultLevel = deriveResultLevel(op, operandStates, &solverRef);
   if (resultLevel.isInt() && resultLevel.getInt() > levelBudget) {
     resultLevel = LevelState(Invalid{});
   }
+  if (resultLevel.isInvalid() && operandsValid) {
+    LLVM_DEBUG({
+      llvm::dbgs() << "LevelAnalysis: Op " << *op
+                   << " became Invalid! Operands: ";
+      for (auto state : operandStates) {
+        state.print(llvm::dbgs());
+        llvm::dbgs() << ", ";
+      }
+      llvm::dbgs() << "\n";
+      for (Value operand : op->getOperands()) {
+        llvm::dbgs() << "  Operand: " << operand << "\n";
+        if (auto* defOp = operand.getDefiningOp()) {
+          llvm::dbgs() << "    Defined by: " << *defOp << "\n";
+        } else {
+          llvm::dbgs() << "    Block argument\n";
+        }
+      }
+    });
+  }
   for (auto result : op->getOpResults()) {
     if (isa<mgmt::InitOp>(op) || isSecretInternal(op, result).value_or(false)) {
       propagate(result, resultLevel);

lib/Analysis/LevelAnalysis/LevelAnalysis.h

@@ -197,6 +197,7 @@ class LevelAnalysis
  public:
   LevelAnalysis(DataFlowSolver& solver, int levelBudget = -1)
       : dataflow::SparseForwardDataFlowAnalysis<LevelLattice>(solver),
+        solverRef(solver),
         levelBudget(levelBudget >= 0 ? levelBudget : kDefaultLevelBudget) {}
   friend class SecretnessAnalysisDependent<LevelAnalysis>;
 
@@ -217,10 +218,12 @@ class LevelAnalysis
   }
 
  private:
+  DataFlowSolver& solverRef;
   int levelBudget;
 };
 
-LevelState deriveResultLevel(Operation* op, ArrayRef<LevelState> operands);
+LevelState deriveResultLevel(Operation* op, ArrayRef<LevelState> operands,
+                             const DataFlowSolver* solver);
 
 /// Backward Analyze the level of plaintext operands of ct-pt ops.
 ///

lib/Dialect/BUILD

@@ -64,6 +64,19 @@ cc_library(
     alwayslink = 1,
 )
 
+cc_library(
+    name = "reduces_level_op_interface_registration",
+    srcs = ["ReducesLevelOpInterfaceRegistration.cpp"],
+    hdrs = ["ReducesLevelOpInterfaceRegistration.h"],
+    deps = [
+        ":HEIRInterfaces",
+        "@heir//lib/Analysis/SecretnessAnalysis",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:LinalgDialect",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
 td_library(
     name = "td_files",
     srcs = ["HEIRInterfaces.td"],

lib/Dialect/HEIRInterfaces.h

@@ -16,6 +16,8 @@
 // IWYU pragma: end_keep
 
 namespace mlir {
+class DataFlowSolver;
+
 namespace heir {
 
 class ElementwiseByOperandOpInterface;

lib/Dialect/HEIRInterfaces.td

@@ -80,12 +80,12 @@ def ReducesLevelOpInterface : OpInterface<"ReducesLevelOpInterface"> {
       /*defaultBody=*/[{ return 1; }]
     >,
     InterfaceMethod<
-      /*desc=*/"Return the OpOperand whose level is reduced.",
-      /*retTy=*/"::mlir::OpOperand&",
-      /*methodName=*/"getOperandToReduce",
-      /*args=*/(ins ),
+      /*desc=*/"Return the OpOperands whose level is reduced.",
+      /*retTy=*/"::llvm::SmallVector<::mlir::OpOperand*>",
+      /*methodName=*/"getOperandsToReduce",
+      /*args=*/(ins "const ::mlir::DataFlowSolver *":$solver),
       /*body=*/[{}],
-      /*defaultBody=*/[{ return $_op->getOpOperand(0); }]
+      /*defaultBody=*/[{ return {&$_op->getOpOperand(0)}; }]
     >
   ];
 }

lib/Dialect/Kernel/IR/KernelOps.cpp

@@ -7,6 +7,7 @@
 #include <utility>
 #include <vector>
 
+#include "lib/Dialect/HEIRInterfaces.h"
 #include "lib/Dialect/LWE/IR/LWEAttributes.h"
 #include "lib/Dialect/LWE/IR/LWETypes.h"
 #include "lib/Target/CompilationTarget/CompilationTarget.h"
@@ -135,14 +136,16 @@ int EvalChebyshevOp::getLevelsToDrop() {
   return baseDepth;
 }
 
-::mlir::OpOperand& EvalChebyshevOp::getOperandToReduce() {
-  return getOperation()->getOpOperand(0);
+::llvm::SmallVector<::mlir::OpOperand*> EvalChebyshevOp::getOperandsToReduce(
+    const ::mlir::DataFlowSolver* solver) {
+  return {&getOperation()->getOpOperand(0)};
 }
 
 int LinearTransformOp::getLevelsToDrop() { return 1; }
 
-::mlir::OpOperand& LinearTransformOp::getOperandToReduce() {
-  return getOperation()->getOpOperand(0);
+::llvm::SmallVector<::mlir::OpOperand*> LinearTransformOp::getOperandsToReduce(
+    const ::mlir::DataFlowSolver* solver) {
+  return {&getOperation()->getOpOperand(0)};
 }
 
 LogicalResult LinearTransformOp::verify() {

lib/Dialect/Kernel/IR/KernelOps.td

@@ -16,7 +16,7 @@ class Kernel_Op<string mnemonic, list<Trait> traits = []> :
 
 def Kernel_EvalChebyshevOp : Kernel_Op<"eval_chebyshev", [
   Pure,
-  DeclareOpInterfaceMethods<ReducesLevelOpInterface, ["getLevelsToDrop", "getOperandToReduce"]>,
+  DeclareOpInterfaceMethods<ReducesLevelOpInterface, ["getLevelsToDrop", "getOperandsToReduce"]>,
   ElementwiseMappable,
   IncreasesMulDepthOpInterface
 ]> {
@@ -40,7 +40,7 @@ def Kernel_EvalChebyshevOp : Kernel_Op<"eval_chebyshev", [
 
 def Kernel_LinearTransformOp : Kernel_Op<"linear_transform", [
   Pure,
-  DeclareOpInterfaceMethods<ReducesLevelOpInterface, ["getLevelsToDrop", "getOperandToReduce"]>,
+  DeclareOpInterfaceMethods<ReducesLevelOpInterface, ["getLevelsToDrop", "getOperandsToReduce"]>,
   ElementwiseMappable,
   IncreasesMulDepthOpInterface
 ]> {

lib/Dialect/Lattigo/IR/LattigoBGVOps.td

@@ -202,7 +202,7 @@ def Lattigo_BGVRelinearizeNewOp : Lattigo_BGVUnaryOp<"relinearize_new"> {
 }
 
 def Lattigo_BGVRescaleNewOp : Lattigo_BGVUnaryOp<"rescale_new", [
-  DeclareOpInterfaceMethods<ReducesLevelOpInterface, ["getOperandToReduce"]>
+  DeclareOpInterfaceMethods<ReducesLevelOpInterface, ["getOperandsToReduce"]>
 ]> {
   let summary = "Rescale a ciphertext in the Lattigo BGV dialect";
   let description = [{
@@ -264,7 +264,7 @@ def Lattigo_BGVRelinearizeOp : Lattigo_BGVUnaryInPlaceOp<"relinearize"> {
 }
 
 def Lattigo_BGVRescaleOp : Lattigo_BGVUnaryInPlaceOp<"rescale", [
-  DeclareOpInterfaceMethods<ReducesLevelOpInterface, ["getOperandToReduce"]>
+  DeclareOpInterfaceMethods<ReducesLevelOpInterface, ["getOperandsToReduce"]>
 ]> {
   let summary = "Rescale a ciphertext in the Lattigo BGV dialect";
   let description = [{

lib/Dialect/Lattigo/IR/LattigoCKKSOps.td

@@ -236,7 +236,7 @@ def Lattigo_CKKSRelinearizeNewOp : Lattigo_CKKSUnaryOp<"relinearize_new"> {
 }
 
 def Lattigo_CKKSRescaleNewOp : Lattigo_CKKSUnaryOp<"rescale_new", [
-  DeclareOpInterfaceMethods<ReducesLevelOpInterface, ["getOperandToReduce"]>
+  DeclareOpInterfaceMethods<ReducesLevelOpInterface, ["getOperandsToReduce"]>
 ]> {
   let summary = "Rescale a ciphertext in the Lattigo CKKS dialect";
   let description = [{
@@ -292,7 +292,7 @@ def Lattigo_CKKSRelinearizeOp : Lattigo_CKKSUnaryInPlaceOp<"relinearize"> {
 }
 
 def Lattigo_CKKSRescaleOp : Lattigo_CKKSUnaryInPlaceOp<"rescale", [
-  DeclareOpInterfaceMethods<ReducesLevelOpInterface, ["getOperandToReduce"]>
+  DeclareOpInterfaceMethods<ReducesLevelOpInterface, ["getOperandsToReduce"]>
 ]> {
   let summary = "Rescale a ciphertext in the Lattigo CKKS dialect";
   let description = [{

lib/Dialect/Lattigo/IR/LattigoOps.cpp

@@ -2,6 +2,7 @@
 
 #include <cstdint>
 
+#include "lib/Dialect/HEIRInterfaces.h"
 #include "lib/Dialect/Lattigo/IR/LattigoTypes.h"
 #include "lib/Utils/RotationUtils.h"
 #include "lib/Utils/Utils.h"
@@ -54,32 +55,38 @@ int RLWEDropLevelNewOp::getLevelsToDrop() { return getLevelToDrop(); }
 
 int RLWEDropLevelOp::getLevelsToDrop() { return getLevelToDrop(); }
 
-::mlir::OpOperand& BGVRescaleNewOp::getOperandToReduce() {
-  return getOperation()->getOpOperand(1);
+::llvm::SmallVector<::mlir::OpOperand*> BGVRescaleNewOp::getOperandsToReduce(
+    const ::mlir::DataFlowSolver* solver) {
+  return {&getOperation()->getOpOperand(1)};
 }
 
-::mlir::OpOperand& BGVRescaleOp::getOperandToReduce() {
-  return getOperation()->getOpOperand(1);
+::llvm::SmallVector<::mlir::OpOperand*> BGVRescaleOp::getOperandsToReduce(
+    const ::mlir::DataFlowSolver* solver) {
+  return {&getOperation()->getOpOperand(1)};
 }
 
-::mlir::OpOperand& CKKSRescaleNewOp::getOperandToReduce() {
-  return getOperation()->getOpOperand(1);
+::llvm::SmallVector<::mlir::OpOperand*> CKKSRescaleNewOp::getOperandsToReduce(
+    const ::mlir::DataFlowSolver* solver) {
+  return {&getOperation()->getOpOperand(1)};
 }
 
-::mlir::OpOperand& CKKSRescaleOp::getOperandToReduce() {
-  return getOperation()->getOpOperand(1);
+::llvm::SmallVector<::mlir::OpOperand*> CKKSRescaleOp::getOperandsToReduce(
+    const ::mlir::DataFlowSolver* solver) {
+  return {&getOperation()->getOpOperand(1)};
 }
 
 ::mlir::OpOperand& CKKSBootstrapOp::getOperandToReset() {
   return getOperation()->getOpOperand(1);
 }
 
-::mlir::OpOperand& RLWEDropLevelNewOp::getOperandToReduce() {
-  return getOperation()->getOpOperand(1);
+::llvm::SmallVector<::mlir::OpOperand*> RLWEDropLevelNewOp::getOperandsToReduce(
+    const ::mlir::DataFlowSolver* solver) {
+  return {&getOperation()->getOpOperand(1)};
 }
 
-::mlir::OpOperand& RLWEDropLevelOp::getOperandToReduce() {
-  return getOperation()->getOpOperand(1);
+::llvm::SmallVector<::mlir::OpOperand*> RLWEDropLevelOp::getOperandsToReduce(
+    const ::mlir::DataFlowSolver* solver) {
+  return {&getOperation()->getOpOperand(1)};
 }
 
 LogicalResult BGVRotateColumnsNewOp::verify() {

lib/Dialect/Lattigo/IR/LattigoRLWEOps.td

@@ -123,7 +123,7 @@ def Lattigo_RLWEDecryptOp : Lattigo_RLWEOp<"decrypt", [Pure]> {
 }
 
 def Lattigo_RLWEDropLevelNewOp : Lattigo_RLWEOp<"drop_level_new",
-  [Pure, DeclareOpInterfaceMethods<ReducesLevelOpInterface, ["getLevelsToDrop", "getOperandToReduce"]>]> {
+  [Pure, DeclareOpInterfaceMethods<ReducesLevelOpInterface, ["getLevelsToDrop", "getOperandsToReduce"]>]> {
   let summary = "Drop level of a ciphertext";
   let arguments = (ins
     Lattigo_RLWEEvaluator:$evaluator,
@@ -134,7 +134,7 @@ def Lattigo_RLWEDropLevelNewOp : Lattigo_RLWEOp<"drop_level_new",
 }
 
 def Lattigo_RLWEDropLevelOp : Lattigo_RLWEOp<"drop_level",
-  [InPlaceOpInterface, DeclareOpInterfaceMethods<ReducesLevelOpInterface, ["getLevelsToDrop", "getOperandToReduce"]>]> {
+  [InPlaceOpInterface, DeclareOpInterfaceMethods<ReducesLevelOpInterface, ["getLevelsToDrop", "getOperandsToReduce"]>]> {
   let summary = "Drop level of a ciphertext";
   let description = [{
     This operation drops the level of a ciphertext