diff --git a/lib/Analysis/RotationAnalysis/BUILD b/lib/Analysis/RotationAnalysis/BUILD index b5d3b563a4..0efde78cac 100644 --- a/lib/Analysis/RotationAnalysis/BUILD +++ b/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", diff --git a/lib/Analysis/RotationAnalysis/DagBuilder.cpp b/lib/Analysis/RotationAnalysis/DagBuilder.cpp index cb71c8f210..ea33084810 100644 --- a/lib/Analysis/RotationAnalysis/DagBuilder.cpp +++ b/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 DagBuilder::visitBlockWithSingleTerminator(Block* block) { return last; } +FailureOr 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 inits; + SmallVector 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& iterArgs) { + valueToNode[op.getInductionVar()] = inductionVar; + for (const auto& [val, node] : + llvm::zip(op.getRegionIterArgs(), iterArgs)) { + valueToNode[val] = node; + } + + FailureOr 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 DagBuilder::visit(affine::AffineYieldOp op) { + std::vector operands; + operands.reserve(op->getNumOperands()); + for (Value operand : op->getOperands()) { + operands.push_back(findNodeOrMakeNewVariable(operand)); + } + return Node::yield(operands); +} + FailureOr DagBuilder::visit(scf::ForOp op) { IntegerAttr lb, ub, step; if (!matchPattern(op.getLowerBound(), m_Constant(&lb)) || @@ -394,9 +445,10 @@ FailureOr DagBuilder::build(Operation* op) { LDBG() << "Visiting op " << *op; return llvm::TypeSwitch>(op) - .Case( [&](auto op) { return visit(op); }) diff --git a/lib/Analysis/RotationAnalysis/DagBuilder.h b/lib/Analysis/RotationAnalysis/DagBuilder.h index a8be741fb5..401b30314b 100644 --- a/lib/Analysis/RotationAnalysis/DagBuilder.h +++ b/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 visit(arith::SubFOp op); FailureOr visit(arith::SubIOp op); FailureOr visit(arith::NegFOp op); + FailureOr visit(affine::AffineForOp op); + FailureOr visit(affine::AffineYieldOp op); FailureOr visit(scf::ForOp op); FailureOr visit(scf::IfOp op); FailureOr visit(scf::YieldOp op); diff --git a/lib/Analysis/RotationAnalysis/RotationEvalVisitorTest.cpp b/lib/Analysis/RotationAnalysis/RotationEvalVisitorTest.cpp index 2222647b7f..5e283cf1c0 100644 --- a/lib/Analysis/RotationAnalysis/RotationEvalVisitorTest.cpp +++ b/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 {}; TEST_P(RollUnrollTest, RotateAndReduceKernel) { diff --git a/tests/Transforms/rotation_analysis/affine_for.mlir b/tests/Transforms/rotation_analysis/affine_for.mlir new file mode 100644 index 0000000000..7b89d63f50 --- /dev/null +++ b/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 +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> + } +}