RotationAnalysis: model nested affine.for in arithmetic DAGs

lib/Analysis/RotationAnalysis/BUILD

@@ -36,6 +36,7 @@ cc_library(
         "@heir//lib/Kernel:ArithmeticDag",
         "@heir//lib/Kernel:Utils",
         "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AffineDialect",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:SCFDialect",

lib/Analysis/RotationAnalysis/DagBuilder.cpp

@@ -8,10 +8,11 @@
 #include "lib/Kernel/AbstractValue.h"
 #include "lib/Kernel/ArithmeticDag.h"
 #include "lib/Kernel/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 "llvm/include/llvm/Support/DebugLog.h"          // from @llvm-project
+#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 "llvm/include/llvm/Support/DebugLog.h"  // from @llvm-project
+#include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"    // from @llvm-project
 #include "mlir/include/mlir/Dialect/SCF/IR/SCF.h"        // from @llvm-project
 #include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
@@ -95,6 +96,56 @@ FailureOr<NodePtr> DagBuilder::visitBlockWithSingleTerminator(Block* block) {
   return last;
 }
 
+FailureOr<NodePtr> DagBuilder::visit(affine::AffineForOp op) {
+  // affine.for carries its bounds as affine maps rather than SSA values, but
+  // is otherwise the same shape as scf.for. Loops whose trip count is not a
+  // compile-time constant cannot be enumerated, same as for scf.for.
+  if (!op.hasConstantBounds()) {
+    LDBG() << "Loop bounds must be constant for RotationAnalysis";
+    return failure();
+  }
+
+  SmallVector<NodePtr> inits;
+  SmallVector<DagType> initTypes;
+  inits.reserve(op.getInits().size());
+  for (Value init : op.getInits()) {
+    NodePtr var = findNodeOrMakeNewVariable(init);
+    valueToNode[init] = var;
+    inits.push_back(var);
+    initTypes.push_back(mlirTypeToDagType(init.getType()));
+  }
+
+  auto dagNode = Node::loop(
+      inits, initTypes, op.getConstantLowerBound(), op.getConstantUpperBound(),
+      op.getStepAsInt(),
+      [&](NodePtr inductionVar, const std::vector<NodePtr>& iterArgs) {
+        valueToNode[op.getInductionVar()] = inductionVar;
+        for (const auto& [val, node] :
+             llvm::zip(op.getRegionIterArgs(), iterArgs)) {
+          valueToNode[val] = node;
+        }
+
+        FailureOr<NodePtr> bodyRes =
+            visitBlockWithSingleTerminator(op.getBody());
+        assert(succeeded(bodyRes) && "failed to parse body");
+        return *bodyRes;
+      });
+
+  for (OpResult opResult : op->getOpResults()) {
+    valueToNode[opResult] = Node::resultAt(dagNode, opResult.getResultNumber());
+  }
+  return dagNode;
+}
+
+FailureOr<NodePtr> DagBuilder::visit(affine::AffineYieldOp op) {
+  std::vector<NodePtr> operands;
+  operands.reserve(op->getNumOperands());
+  for (Value operand : op->getOperands()) {
+    operands.push_back(findNodeOrMakeNewVariable(operand));
+  }
+  return Node::yield(operands);
+}
+
 FailureOr<NodePtr> DagBuilder::visit(scf::ForOp op) {
   IntegerAttr lb, ub, step;
   if (!matchPattern(op.getLowerBound(), m_Constant(&lb)) ||
@@ -394,9 +445,10 @@ FailureOr<NodePtr> DagBuilder::build(Operation* op) {
   LDBG() << "Visiting op " << *op;
 
   return llvm::TypeSwitch<Operation*, FailureOr<NodePtr>>(op)
-      .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,
+      .Case<affine::AffineForOp, affine::AffineYieldOp, 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::InsertOp,
             tensor::InsertSliceOp, tensor::SplatOp, RotationOpInterface>(
           [&](auto op) { return visit(op); })

lib/Analysis/RotationAnalysis/DagBuilder.h

@@ -6,6 +6,7 @@
 #include "lib/Dialect/HEIRInterfaces.h"
 #include "lib/Kernel/AbstractValue.h"
 #include "lib/Kernel/ArithmeticDag.h"
+#include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"    // from @llvm-project
 #include "mlir/include/mlir/Dialect/SCF/IR/SCF.h"        // from @llvm-project
 #include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
@@ -45,6 +46,8 @@ class DagBuilder {
   FailureOr<NodePtr> visit(arith::SubFOp op);
   FailureOr<NodePtr> visit(arith::SubIOp op);
   FailureOr<NodePtr> visit(arith::NegFOp op);
+  FailureOr<NodePtr> visit(affine::AffineForOp op);
+  FailureOr<NodePtr> visit(affine::AffineYieldOp op);
   FailureOr<NodePtr> visit(scf::ForOp op);
   FailureOr<NodePtr> visit(scf::IfOp op);
   FailureOr<NodePtr> visit(scf::YieldOp op);

lib/Analysis/RotationAnalysis/RotationEvalVisitorTest.cpp

@@ -36,6 +36,27 @@ TEST(RotationEvalVisitorTest, TestSimpleLoop) {
   EXPECT_THAT(evalRotations(loop), UnorderedElementsAre(0, 2, 4, 6, 8));
 }
 
+// A loop that rotates by both the induction variable and by
+// inductionVar + constant. This is the shape a rolled convolution emits, and
+// the offset rotations are the ones whose Galois keys went missing (a rotation
+// by 65 = 1 + 64 aborted lattigo with "GaloisKey[39685] is nil").
+TEST(RotationEvalVisitorTest, TestLoopWithOffsetRotation) {
+  LiteralValue inputVector({1, 2, 3, 4, 5, 6, 7, 8, 9, 10});
+  auto x = Node::leaf(inputVector);
+  auto four = Node::constantScalar(4, DagType::index());
+  auto loop = Node::loop(x, {DagType::intTensor(32, {10})}, 1, 6, 1,
+                         [&](NodePtr iv, NodePtr iterArg) {
+                           NodePtr plain = Node::leftRotate(iterArg, iv);
+                           NodePtr offset = Node::add(iv, four);
+                           NodePtr shifted = Node::leftRotate(plain, offset);
+                           return Node::yield({shifted});
+                         });
+  // iv runs 1..5, so the plain rotations are 1..5 and the offset rotations are
+  // 5..9; normalized modulo the 10-slot dimension they stay distinct.
+  EXPECT_THAT(evalRotations(loop),
+              UnorderedElementsAre(1, 2, 3, 4, 5, 6, 7, 8, 9));
+}
+
 class RollUnrollTest : public testing::TestWithParam<bool> {};
 
 TEST_P(RollUnrollTest, RotateAndReduceKernel) {

tests/Transforms/rotation_analysis/affine_for.mlir

@@ -0,0 +1,33 @@
+// RUN: heir-opt --rotation-analysis --split-input-file %s | FileCheck %s
+
+// The loop's yielded value is produced by an affine.for nested inside the
+// scf.for. Unless affine.for is modelled, its result is resolved as an opaque
+// variable and every rotation feeding it is severed from the DAG root, so
+// handleScfFor reports success having found no shifts at all. Both the bare
+// induction-variable rotation and the inductionVar + constant one must be
+// recovered: a rolled convolution emits the pair, and the offset one is what
+// aborted lattigo with "GaloisKey[39685] is nil" (39685 = 5^65).
+
+// CHECK: module attributes
+// CHECK-SAME: rotation_analysis.indices = array<i64: 1, 2, 3, 4, 5, 65, 66, 67, 68, 69>
+module attributes {scheme.actual_slot_count = 128} {
+  func.func @rotations_yielded_through_affine_for(
+      %arg0: tensor<128xi32>) -> tensor<128xi32> {
+    %c1 = arith.constant 1 : index
+    %c6 = arith.constant 6 : index
+    %c64 = arith.constant 64 : index
+    %0 = scf.for %i = %c1 to %c6 step %c1 iter_args(%iter = %arg0)
+        -> (tensor<128xi32>) {
+      %r0 = tensor_ext.rotate %arg0, %i : tensor<128xi32>, index
+      %off = arith.addi %i, %c64 : index
+      %r1 = tensor_ext.rotate %arg0, %off : tensor<128xi32>, index
+      %sum = arith.addi %r0, %r1 : tensor<128xi32>
+      %1 = affine.for %j = 0 to 2 iter_args(%acc = %sum) -> (tensor<128xi32>) {
+        %2 = arith.addi %acc, %sum : tensor<128xi32>
+        affine.yield %2 : tensor<128xi32>
+      }
+      scf.yield %1 : tensor<128xi32>
+    }
+    return %0 : tensor<128xi32>
+  }
+}