Make rotation analysis more robust for loops with if statements

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) {