Make rotation analysis more robust for loops with if statements

Here the problem was that, for more complicated looped kernels,
the rotation analysis was not evaluating the loop correctly.
The previous implementation was trying to skip various intermediate
operations on tensors/ciphertexts, but this causes the DAG to be disconnected
and so the loop's yield may not end up being visited, so that it doesn't
proceed to the next iteration and truncates the rotations at the first
one found (without these changes, the if_else.mlir test will only annotate with
the index 512).

PiperOrigin-RevId: 885851705

Author: j2kun

lib/Analysis/RotationAnalysis/BUILD

@@ -30,6 +30,7 @@ cc_library(
     srcs = ["DagBuilder.cpp"],
     hdrs = ["DagBuilder.h"],
     deps = [
+        "@heir//lib/Dialect:HEIRInterfaces",
         "@heir//lib/Dialect/TensorExt/IR:Dialect",
         "@heir//lib/Kernel:AbstractValue",
         "@heir//lib/Kernel:ArithmeticDag",

lib/Analysis/RotationAnalysis/DagBuilder.cpp

@@ -4,6 +4,7 @@
 #include <memory>
 #include <vector>
 
+#include "lib/Dialect/HEIRInterfaces.h"
 #include "lib/Dialect/TensorExt/IR/TensorExtOps.h"
 #include "lib/Kernel/AbstractValue.h"
 #include "lib/Kernel/ArithmeticDag.h"
@@ -135,6 +136,29 @@ FailureOr<NodePtr> DagBuilder::visit(scf::ForOp op) {
   return dagNode;
 }
 
+FailureOr<NodePtr> DagBuilder::visit(scf::IfOp op) {
+  if (op.getElseRegion().empty()) {
+    op.emitOpError()
+        << "skipped: If/Else without an else branch is not supported\n";
+    return failure();
+  }
+
+  auto condition = findNodeOrMakeNewVariable(op.getCondition());
+  auto thenRes = visitBlockWithSingleTerminator(&op.getThenRegion().front());
+  if (failed(thenRes)) return failure();
+  NodePtr thenBody = thenRes.value();
+
+  auto elseRes = visitBlockWithSingleTerminator(&op.getElseRegion().front());
+  if (failed(elseRes)) return failure();
+  NodePtr elseBody = elseRes.value();
+
+  auto dagNode = Node::ifElse(condition, thenBody, elseBody);
+  for (OpResult opResult : op->getOpResults()) {
+    valueToNode[opResult] = Node::resultAt(dagNode, opResult.getResultNumber());
+  }
+  return dagNode;
+}
+
 FailureOr<NodePtr> DagBuilder::visit(scf::YieldOp op) {
   std::vector<NodePtr> operands;
   operands.reserve(op->getNumOperands());
@@ -144,11 +168,32 @@ FailureOr<NodePtr> DagBuilder::visit(scf::YieldOp op) {
   return Node::yield(operands);
 }
 
-FailureOr<NodePtr> DagBuilder::visit(tensor_ext::RotateOp op) {
-  auto tensor = findNodeOrMakeNewVariable(op.getTensor());
-  auto shift = findNodeOrMakeNewVariable(op.getShift());
-  auto dagNode = Node::leftRotate(tensor, shift);
-  valueToNode[op.getResult()] = dagNode;
+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));
+  }
+
+  // Find the rotatable operand.
+  // We assume it's the operand that has the same type as the result.
+  OpOperand* rotatedOperand = op.getRotatedOperand();
+
+  if (!rotatedOperand) {
+    LDBG() << "Could not find a rotated operand for op " << op
+           << ". It may be that the default implementation of "
+              "getRotatedOperand is incorrect for this RotationOpInterface.";
+    return failure();
+  }
+
+  auto tensorNode = findNodeOrMakeNewVariable(rotatedOperand->get());
+  auto dagNode = Node::leftRotate(tensorNode, shift);
+  valueToNode[op->getResult(0)] = dagNode;
   return dagNode;
 }
 
@@ -176,6 +221,41 @@ FailureOr<NodePtr> DagBuilder::visit(arith::SubIOp op) {
   return dagNode;
 }
 
+FailureOr<NodePtr> DagBuilder::visit(arith::CmpIOp op) {
+  auto lhs = findNodeOrMakeNewVariable(op.getLhs());
+  auto rhs = findNodeOrMakeNewVariable(op.getRhs());
+
+  kernel::ComparisonPredicate pred;
+  switch (op.getPredicate()) {
+    case arith::CmpIPredicate::slt:
+    case arith::CmpIPredicate::ult:
+      pred = kernel::ComparisonPredicate::LT;
+      break;
+    case arith::CmpIPredicate::sle:
+    case arith::CmpIPredicate::ule:
+      pred = kernel::ComparisonPredicate::LE;
+      break;
+    case arith::CmpIPredicate::sgt:
+    case arith::CmpIPredicate::ugt:
+      pred = kernel::ComparisonPredicate::GT;
+      break;
+    case arith::CmpIPredicate::sge:
+    case arith::CmpIPredicate::uge:
+      pred = kernel::ComparisonPredicate::GE;
+      break;
+    case arith::CmpIPredicate::eq:
+      pred = kernel::ComparisonPredicate::EQ;
+      break;
+    case arith::CmpIPredicate::ne:
+      pred = kernel::ComparisonPredicate::NE;
+      break;
+  }
+
+  auto dagNode = Node::comparison(lhs, rhs, pred);
+  valueToNode[op.getResult()] = dagNode;
+  return dagNode;
+}
+
 FailureOr<NodePtr> DagBuilder::visit(arith::ConstantOp op) {
   NodePtr dagNode =
       TypeSwitch<Attribute, NodePtr>(op.getValue())
@@ -234,16 +314,61 @@ FailureOr<NodePtr> DagBuilder::visit(arith::DivSIOp op) {
   return dagNode;
 }
 
+FailureOr<NodePtr> DagBuilder::visit(arith::AddFOp op) {
+  auto lhs = findNodeOrMakeNewVariable(op.getLhs());
+  auto rhs = findNodeOrMakeNewVariable(op.getRhs());
+  auto dagNode = Node::add(lhs, rhs);
+  valueToNode[op.getResult()] = dagNode;
+  return dagNode;
+}
+
+FailureOr<NodePtr> DagBuilder::visit(arith::MulFOp op) {
+  auto lhs = findNodeOrMakeNewVariable(op.getLhs());
+  auto rhs = findNodeOrMakeNewVariable(op.getRhs());
+  auto dagNode = Node::mul(lhs, rhs);
+  valueToNode[op.getResult()] = dagNode;
+  return dagNode;
+}
+
+FailureOr<NodePtr> DagBuilder::visit(arith::SubFOp op) {
+  auto lhs = findNodeOrMakeNewVariable(op.getLhs());
+  auto rhs = findNodeOrMakeNewVariable(op.getRhs());
+  auto dagNode = Node::sub(lhs, rhs);
+  valueToNode[op.getResult()] = dagNode;
+  return dagNode;
+}
+
+FailureOr<NodePtr> DagBuilder::visit(arith::NegFOp op) {
+  auto lhs = findNodeOrMakeNewVariable(op.getOperand());
+  // ArithmeticDag doesn't have NegNode, so we can use (0 - lhs)
+  auto zero = Node::constantScalar(0.0, mlirTypeToDagType(op.getType()));
+  auto dagNode = Node::sub(zero, lhs);
+  valueToNode[op.getResult()] = dagNode;
+  return dagNode;
+}
+
+FailureOr<NodePtr> DagBuilder::visit(tensor::ExtractSliceOp op) {
+  // For rotation analysis, we can approximate extract_slice as just returning
+  // the source tensor if we don't care about the specific values, or we can
+  // try to be more precise. Since rotation analysis usually only cares about
+  // the fact that *some* tensor is being rotated, returning the source is
+  // enough to keep the DAG connected.
+  auto source = findNodeOrMakeNewVariable(op.getSource());
+  valueToNode[op.getResult()] = source;
+  return source;
+}
+
 FailureOr<NodePtr> DagBuilder::build(Operation* op) {
   LDBG() << "Visiting op " << *op;
   return llvm::TypeSwitch<Operation*, FailureOr<NodePtr>>(op)
-      .Case<arith::AddIOp, arith::ConstantOp, arith::DivSIOp, arith::MulIOp,
-            arith::SubIOp, scf::ForOp, scf::YieldOp, tensor::ExtractOp,
-            tensor::SplatOp, tensor_ext::RotateOp>(
-          [&](auto op) { return visit(op); })
-      .Default([&](auto op) {
-        LDBG() << "Unsupported op type " << op->getName();
-        return failure();
+      .Case<arith::AddFOp, arith::AddIOp, arith::CmpIOp, arith::ConstantOp,
+            arith::DivSIOp, arith::MulFOp, arith::MulIOp, arith::SubFOp,
+            arith::SubIOp, arith::NegFOp, scf::ForOp, scf::IfOp, scf::YieldOp,
+            tensor::ExtractOp, tensor::ExtractSliceOp, tensor::SplatOp,
+            RotationOpInterface>([&](auto op) { return visit(op); })
+      .Default([&](Operation* op) -> FailureOr<NodePtr> {
+        LDBG() << "Unsupported op type " << op->getName() << ", skipping";
+        return NodePtr(nullptr);
       });
 }
 

lib/Analysis/RotationAnalysis/DagBuilder.h

@@ -3,6 +3,7 @@
 
 #include <memory>
 
+#include "lib/Dialect/HEIRInterfaces.h"
 #include "lib/Dialect/TensorExt/IR/TensorExtOps.h"
 #include "lib/Kernel/AbstractValue.h"
 #include "lib/Kernel/ArithmeticDag.h"
@@ -35,16 +36,23 @@ class DagBuilder {
   // Visit a block and return the NodePtr corresponding to its terminator
   FailureOr<NodePtr> visitBlockWithSingleTerminator(Block* block);
 
+  FailureOr<NodePtr> visit(arith::AddFOp op);
   FailureOr<NodePtr> visit(arith::AddIOp op);
+  FailureOr<NodePtr> visit(arith::CmpIOp op);
   FailureOr<NodePtr> visit(arith::ConstantOp op);
   FailureOr<NodePtr> visit(arith::DivSIOp op);
+  FailureOr<NodePtr> visit(arith::MulFOp op);
   FailureOr<NodePtr> visit(arith::MulIOp op);
+  FailureOr<NodePtr> visit(arith::SubFOp op);
   FailureOr<NodePtr> visit(arith::SubIOp op);
+  FailureOr<NodePtr> visit(arith::NegFOp op);
   FailureOr<NodePtr> visit(scf::ForOp op);
+  FailureOr<NodePtr> visit(scf::IfOp op);
   FailureOr<NodePtr> visit(scf::YieldOp op);
   FailureOr<NodePtr> visit(tensor::ExtractOp op);
+  FailureOr<NodePtr> visit(tensor::ExtractSliceOp op);
   FailureOr<NodePtr> visit(tensor::SplatOp op);
-  FailureOr<NodePtr> visit(tensor_ext::RotateOp op);
+  FailureOr<NodePtr> visit(RotationOpInterface op);
 
   // A mapping of previously visited Values
   DenseMap<Value, NodePtr> valueToNode;

lib/Analysis/RotationAnalysis/RotationAnalysis.cpp

@@ -47,7 +47,7 @@ LogicalResult RotationAnalysis::handleScfFor(scf::ForOp forOp) {
 
   NodePtr dag = res.value();
   auto shifts = evalRotations(dag);
-  rotationIndices = DenseSet<int64_t>(shifts.begin(), shifts.end());
+  rotationIndices.insert(shifts.begin(), shifts.end());
 
   // All the rotation ops within the outermost for loop are analyzed.
   outermostFor->walk([&](RotationOpInterface rotOp) { markVisited(rotOp); });

lib/Analysis/RotationAnalysis/RotationEvalVisitor.cpp

@@ -27,31 +27,24 @@ using kernel::LiteralValue;
 using kernel::VariableNode;
 
 // This is a copy of EvalVisitor::operator() for LeftRotateNode, but recording
-// the materialized rotation amount.
+// the materialized rotation amount and skipping the actual computation.
 EvalResults RotationEvalVisitor::operator()(
     const LeftRotateNode<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[0];
+  auto dim = shape.back();  // Use the slot dimension
 
   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;
-
-  // Save the evaluated shift
   evaluatedShifts.insert(amount);
 
-  const auto& oVal = operand.get();
-  const auto* oVec = std::get_if<std::vector<int>>(&oVal);
-  assert(oVec && "unsupported rotate operand type");
-
-  std::vector<int> result(dim);
-  for (size_t i = 0; i < dim; ++i) {
-    result[i] = (*oVec)[(i + amount) % oVec->size()];
-  }
-  return {result};
+  // 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()(

lib/Dialect/HEIRInterfaces.td

@@ -451,6 +451,23 @@ def RotationOpInterface : OpInterface<"RotationOpInterface"> {
       "Returns the rotation index as an OpFoldResult (Value or Attribute).",
       "::mlir::OpFoldResult", "getRotationIndex"
     >,
+
+    InterfaceMethod<
+      /*desc=*/"Returns the rotated operand.",
+      /*retTy=*/"::mlir::OpOperand *",
+      /*methodName=*/"getRotatedOperand",
+      /*args=*/(ins ),
+      /*defaultImplementation=*/[{
+        // By default, find the first operand that matches the first result
+        // type.
+        for (::mlir::OpOperand &operand : $_op->getOpOperands()) {
+          if (operand.get().getType() == $_op->getResult(0).getType()) {
+            return &operand;
+          }
+        }
+        return nullptr;
+      }]
+    >,
   ];
 }
 #endif  // LIB_DIALECT_HEIR_IR_HEIRINTERFACES_TD_

lib/Kernel/EvalVisitor.cpp

@@ -57,13 +57,17 @@ EvalResults EvalVisitor::operator()(const AddNode<LiteralValue>& node) {
   auto dim = left.getShape()[0];
   const auto* lVec = std::get_if<std::vector<int>>(&lVal);
   const auto* rVec = std::get_if<std::vector<int>>(&rVal);
-  assert(lVec && rVec && "unsupported add operands");
-  assert(left.getShape() == right.getShape() && "disagreeing shapes");
-  std::vector<int> result(dim);
-  for (size_t i = 0; i < dim; ++i) {
-    result[i] = (*lVec)[i] + (*rVec)[i];
+  if (lVec && rVec) {
+    assert(left.getShape() == right.getShape() && "disagreeing shapes");
+    std::vector<int> result(dim);
+    for (size_t i = 0; i < dim; ++i) {
+      result[i] = (*lVec)[i] + (*rVec)[i];
+    }
+    return {result};
   }
-  return {result};
+
+  // If types are not supported, just return left as dummy
+  return {left};
 }
 
 EvalResults EvalVisitor::operator()(const SubtractNode<LiteralValue>& node) {
@@ -84,14 +88,16 @@ EvalResults EvalVisitor::operator()(const SubtractNode<LiteralValue>& node) {
   auto dim = left.getShape()[0];
   const auto* lVec = std::get_if<std::vector<int>>(&lVal);
   const auto* rVec = std::get_if<std::vector<int>>(&rVal);
-  assert(lVec && rVec && "unsupported sub operands");
-  assert(left.getShape() == right.getShape() && "disagreeing shapes");
-  std::vector<int> result(dim);
-  for (size_t i = 0; i < dim; ++i) {
-    result[i] = (*lVec)[i] - (*rVec)[i];
+  if (lVec && rVec) {
+    assert(left.getShape() == right.getShape() && "disagreeing shapes");
+    std::vector<int> result(dim);
+    for (size_t i = 0; i < dim; ++i) {
+      result[i] = (*lVec)[i] - (*rVec)[i];
+    }
+    return {result};
   }
 
-  return {result};
+  return {left};
 }
 
 EvalResults EvalVisitor::operator()(const MultiplyNode<LiteralValue>& node) {
@@ -112,13 +118,15 @@ EvalResults EvalVisitor::operator()(const MultiplyNode<LiteralValue>& node) {
   auto dim = left.getShape()[0];
   const auto* lVec = std::get_if<std::vector<int>>(&lVal);
   const auto* rVec = std::get_if<std::vector<int>>(&rVal);
-  assert(lVec && rVec && "unsupported mul operands");
-  assert(left.getShape() == right.getShape() && "disagreeing shapes");
-  std::vector<int> result(dim);
-  for (size_t i = 0; i < dim; ++i) {
-    result[i] = (*lVec)[i] * (*rVec)[i];
+  if (lVec && rVec) {
+    assert(left.getShape() == right.getShape() && "disagreeing shapes");
+    std::vector<int> result(dim);
+    for (size_t i = 0; i < dim; ++i) {
+      result[i] = (*lVec)[i] * (*rVec)[i];
+    }
+    return {result};
   }
-  return {result};
+  return {left};
 }
 
 EvalResults EvalVisitor::operator()(const FloorDivNode<LiteralValue>& node) {
@@ -147,20 +155,23 @@ EvalResults EvalVisitor::operator()(const FloorDivNode<LiteralValue>& node) {
 EvalResults EvalVisitor::operator()(const LeftRotateNode<LiteralValue>& node) {
   LDBG() << "Visiting LeftRotateNode";
   auto operand = this->process(node.operand)[0];
-  auto dim = operand.getShape()[0];
+  auto dim = operand.getShape().back();
   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;
 
   const auto& oVal = operand.get();
-  const auto* oVec = std::get_if<std::vector<int>>(&oVal);
-  assert(oVec && "unsupported rotate operand");
-  std::vector<int> result(dim);
-  for (size_t i = 0; i < dim; ++i) {
-    result[i] = (*oVec)[(i + amount) % oVec->size()];
+  if (const auto* oVec = std::get_if<std::vector<int>>(&oVal)) {
+    std::vector<int> result(dim);
+    for (size_t i = 0; i < dim; ++i) {
+      result[i] = (*oVec)[(i + amount) % oVec->size()];
+    }
+    return {result};
   }
-  return {result};
+
+  // If the operand is not a 1D vector (e.g., a 2D float tensor), return as-is.
+  return {operand};
 }
 
 EvalResults EvalVisitor::operator()(const ExtractNode<LiteralValue>& node) {