Merge pull request #2818 from AlexanderViand:analysis-interfaces

PiperOrigin-RevId: 891704982

Author: copybara-github

lib/Analysis/RotationAnalysis/BUILD

@@ -52,6 +52,7 @@ cc_library(
         "@heir//lib/Kernel:AbstractValue",
         "@heir//lib/Kernel:ArithmeticDag",
         "@heir//lib/Kernel:EvalVisitor",
+        "@heir//lib/Utils:RotationUtils",
         "@llvm-project//llvm:Support",
     ],
 )

lib/Analysis/RotationAnalysis/DagBuilder.cpp

@@ -170,18 +170,17 @@ FailureOr<NodePtr> DagBuilder::visit(scf::YieldOp op) {
 
 FailureOr<NodePtr> DagBuilder::visit(RotationOpInterface op) {
   LDBG() << "Processing RotationOpInterface " << op;
-  OpFoldResult ofr = op.getRotationIndex();
-  NodePtr shift;
-  if (auto attr = dyn_cast<Attribute>(ofr)) {
-    auto intAttr = cast<IntegerAttr>(attr);
-    shift = Node::constantScalar(intAttr.getInt(),
-                                 mlirTypeToDagType(intAttr.getType()));
-  } else {
-    shift = findNodeOrMakeNewVariable(cast<Value>(ofr));
-  }
+  auto indices = op.getRotationIndices();
+  auto toShiftNode = [&](OpFoldResult ofr) -> NodePtr {
+    if (auto attr = dyn_cast<Attribute>(ofr)) {
+      auto intAttr = cast<IntegerAttr>(attr);
+      return Node::constantScalar(intAttr.getInt(),
+                                  mlirTypeToDagType(intAttr.getType()));
+    }
+    return findNodeOrMakeNewVariable(cast<Value>(ofr));
+  };
 
   // Find the rotatable operand.
-  // We assume it's the operand that has the same type as the result.
   OpOperand* rotatedOperand = op.getRotatedOperand();
 
   if (!rotatedOperand) {
@@ -192,7 +191,18 @@ FailureOr<NodePtr> DagBuilder::visit(RotationOpInterface op) {
   }
 
   auto tensorNode = findNodeOrMakeNewVariable(rotatedOperand->get());
-  auto dagNode = Node::leftRotate(tensorNode, shift);
+
+  NodePtr dagNode;
+  if (indices.size() == 1) {
+    dagNode = Node::leftRotate(tensorNode, toShiftNode(indices[0]));
+  } else {
+    std::vector<NodePtr> shifts;
+    shifts.reserve(indices.size());
+    for (auto& ofr : indices) {
+      shifts.push_back(toShiftNode(ofr));
+    }
+    dagNode = Node::leftRotateBulk(tensorNode, std::move(shifts));
+  }
   valueToNode[op->getResult(0)] = dagNode;
   return dagNode;
 }

lib/Analysis/RotationAnalysis/RotationAnalysis.cpp

@@ -65,19 +65,28 @@ LogicalResult RotationAnalysis::analyzeRotationOp(
 
   // Handle cases where the rotation shift can be statically folded via constant
   // propagation.
-  OpFoldResult ofr = rotationOp.getRotationIndex();
-  if (auto attr = dyn_cast_if_present<Attribute>(ofr)) {
-    if (auto intAttr = dyn_cast<IntegerAttr>(attr)) {
-      rotationIndices.insert(intAttr.getInt());
-      markVisited(rotationOp);
-      return success();
+  auto indices = rotationOp.getRotationIndices();
+  SmallVector<int64_t> constantIndices;
+  bool allConstant = true;
+  for (auto& ofr : indices) {
+    if (auto attr = dyn_cast_if_present<Attribute>(ofr)) {
+      if (auto intAttr = dyn_cast<IntegerAttr>(attr)) {
+        constantIndices.push_back(intAttr.getInt());
+        continue;
+      }
     }
+    if (auto value = dyn_cast<Value>(ofr)) {
+      IntegerAttr attr;
+      if (matchPattern(value, m_Constant(&attr))) {
+        constantIndices.push_back(attr.getInt());
+        continue;
+      }
+    }
+    allConstant = false;
+    break;
   }
-
-  Value value = dyn_cast<Value>(ofr);
-  IntegerAttr attr;
-  if (matchPattern(value, m_Constant(&attr))) {
-    rotationIndices.insert(attr.getInt());
+  if (allConstant) {
+    rotationIndices.insert(constantIndices.begin(), constantIndices.end());
     markVisited(rotationOp);
     return success();
   }

lib/Analysis/RotationAnalysis/RotationEvalVisitor.cpp

@@ -12,6 +12,7 @@
 #include "lib/Kernel/AbstractValue.h"
 #include "lib/Kernel/ArithmeticDag.h"
 #include "lib/Kernel/EvalVisitor.h"
+#include "lib/Utils/RotationUtils.h"
 #include "llvm/include/llvm/Support/ErrorHandling.h"  // from @llvm-project
 
 #define DEBUG_TYPE "rotation-analysis"
@@ -22,6 +23,7 @@ namespace heir {
 using kernel::ArithmeticDagNode;
 using kernel::DagType;
 using kernel::EvalResults;
+using kernel::LeftRotateBulkNode;
 using kernel::LeftRotateNode;
 using kernel::LiteralValue;
 using kernel::VariableNode;
@@ -38,15 +40,33 @@ EvalResults RotationEvalVisitor::operator()(
   auto evaluatedShift = this->process(node.shift)[0];
   int amount = std::get<int>(evaluatedShift.get());
 
-  // Normalize amount to be in [0, dim)
-  amount = ((amount % dim) + dim) % dim;
+  amount = normalizeRotation(amount, dim);
   evaluatedShifts.insert(amount);
 
   // We don't need to rotate the values for rotation analysis. We just return
   // the operand as-is to keep the IR connected.
   return {operand};
 }
 
+EvalResults RotationEvalVisitor::operator()(
+    const LeftRotateBulkNode<LiteralValue>& node) {
+  auto operand = this->process(node.operand)[0];
+  auto shape = operand.getShape();
+  assert(!shape.empty() && "rotate operand must be a tensor");
+  auto dim = shape.back();
+
+  for (const auto& shiftNode : node.shifts) {
+    auto evaluatedShift = this->process(shiftNode)[0];
+    int amount = std::get<int>(evaluatedShift.get());
+    amount = normalizeRotation(amount, dim);
+    evaluatedShifts.insert(amount);
+  }
+
+  // We don't need to rotate the values for rotation analysis. We just return
+  // the operand as-is to keep the IR connected.
+  return {operand};
+}
+
 EvalResults RotationEvalVisitor::operator()(
     const VariableNode<LiteralValue>& node) {
   if (node.value.has_value()) {

lib/Analysis/RotationAnalysis/RotationEvalVisitor.h

@@ -20,6 +20,10 @@ class RotationEvalVisitor : public kernel::EvalVisitor {
   kernel::EvalResults operator()(
       const kernel::LeftRotateNode<kernel::LiteralValue>& node) override;
 
+  // Override the bulk rotation op to record all materialized rotation shifts.
+  kernel::EvalResults operator()(
+      const kernel::LeftRotateBulkNode<kernel::LiteralValue>& node) override;
+
   // Override the variable node to allow uninitialized values to be populated
   // with anything.
   kernel::EvalResults operator()(

lib/Dialect/BGV/IR/BGVOps.cpp

@@ -36,14 +36,15 @@ LogicalResult ModulusSwitchOp::verify() {
 
 LogicalResult LevelReduceOp::verify() { return lwe::verifyLevelReduceOp(this); }
 
-::mlir::OpFoldResult RotateColumnsOp::getRotationIndex() {
-  if (getStaticShift()) return getStaticShiftAttr();
-  return getDynamicShift();
+::llvm::SmallVector<::mlir::OpFoldResult>
+RotateColumnsOp::getRotationIndices() {
+  if (getStaticShift()) return {getStaticShiftAttr()};
+  return {getDynamicShift()};
 }
 
-::mlir::OpFoldResult RotateRowsOp::getRotationIndex() {
-  if (getStaticShift()) return getStaticShiftAttr();
-  return getDynamicShift();
+::llvm::SmallVector<::mlir::OpFoldResult> RotateRowsOp::getRotationIndices() {
+  if (getStaticShift()) return {getStaticShiftAttr()};
+  return {getDynamicShift()};
 }
 
 //===----------------------------------------------------------------------===//

lib/Dialect/CKKS/IR/CKKSOps.cpp

@@ -47,9 +47,9 @@ LogicalResult RescaleOp::verify() {
 
 LogicalResult LevelReduceOp::verify() { return lwe::verifyLevelReduceOp(this); }
 
-::mlir::OpFoldResult RotateOp::getRotationIndex() {
-  if (getStaticShift()) return getStaticShiftAttr();
-  return getDynamicShift();
+::llvm::SmallVector<::mlir::OpFoldResult> RotateOp::getRotationIndices() {
+  if (getStaticShift()) return {getStaticShiftAttr()};
+  return {getDynamicShift()};
 }
 
 LogicalResult BootstrapOp::verify() {

lib/Dialect/HEIRInterfaces.td

@@ -443,13 +443,15 @@ def OperandAndResultAttrInterface : OpInterface<"OperandAndResultAttrInterface">
 def RotationOpInterface : OpInterface<"RotationOpInterface"> {
   let cppNamespace = "::mlir::heir";
   let description = [{
-    An interface for rotation operations.
+    An interface for operations that perform rotations, possibly as a subroutine
+    of this op. This is used in the analysis step to determine what rotation
+    keys need to be generated.
   }];
 
   let methods = [
     InterfaceMethod<
-      "Returns the rotation index as an OpFoldResult (Value or Attribute).",
-      "::mlir::OpFoldResult", "getRotationIndex"
+      "Returns the rotation indices as OpFoldResults (Values or Attributes).",
+      "::llvm::SmallVector<::mlir::OpFoldResult>", "getRotationIndices"
     >,
 
     InterfaceMethod<

lib/Dialect/Lattigo/IR/BUILD

@@ -90,6 +90,7 @@ cc_library(
         ":types_inc_gen",
         "@heir//lib/Dialect:HEIRInterfaces",
         "@heir//lib/Utils",
+        "@heir//lib/Utils:RotationUtils",
         "@heir//lib/Utils/Tablegen:InPlaceOpInterface",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:InferTypeOpInterface",

lib/Dialect/Lattigo/IR/LattigoCKKSOps.td

@@ -341,7 +341,9 @@ def Lattigo_CKKSBootstrapOp : Lattigo_CKKSUnaryOp<"bootstrap"> {
   let results = (outs Lattigo_RLWECiphertext:$output);
 }
 
-def Lattigo_CKKSLinearTransformOp : Lattigo_CKKSOp<"linear_transform"> {
+def Lattigo_CKKSLinearTransformOp : Lattigo_CKKSOp<"linear_transform", [
+  DeclareOpInterfaceMethods<RotationOpInterface>
+]> {
   let summary = "Apply a linear transform on a lattigo CKKS ciphertext";
   let description = [{
     This operation applies a linear transform on a CKKS ciphertext using