add objectFIFO unrolling pass with constant folding and tests

Author: andrej

include/aie/Dialect/AIE/IR/AIEOps.td

@@ -366,7 +366,7 @@ def AIE_ShimDMAOp: AIE_Op<"shim_dma", [
 }
 
 def AIE_CoreOp: AIE_Op<"core", [
-    IsFlowEndPoint, IsCoreTile, TileElement,
+    IsFlowEndPoint, IsCoreTile, TileElement, AutomaticAllocationScope,
     DeclareOpInterfaceMethods<InferTypeOpInterface>
   ]>, Results<(outs Index)> {
   let arguments = (

include/aie/Dialect/AIE/Transforms/AIEPasses.h

@@ -19,6 +19,13 @@
 
 namespace xilinx::AIE {
 
+/// Discardable attribute set by `AIEObjectFifoStatefulTransform` on each
+/// `scf.for` loop that carries an objectFifo access. Its i64 value is the loop
+/// unroll factor (the least common multiple of the depths of the objectFifos
+/// accessed within the loop) consumed by the `AIEObjectFifoUnroll` pass.
+inline constexpr llvm::StringLiteral kObjectFifoUnrollHintAttrName =
+    "aie.unroll_hint";
+
 #define GEN_PASS_DECL
 #define GEN_PASS_DEF_AIEROUTEPATHFINDERFLOWS
 #include "aie/Dialect/AIE/Transforms/AIEPasses.h.inc"
@@ -54,6 +61,7 @@ createAIEHoistVectorTransferPointersPass();
 std::unique_ptr<mlir::OperationPass<DeviceOp>> createAIEPathfinderPass();
 std::unique_ptr<mlir::OperationPass<DeviceOp>>
 createAIEObjectFifoStatefulTransformPass();
+std::unique_ptr<mlir::OperationPass<DeviceOp>> createAIEObjectFifoUnrollPass();
 std::unique_ptr<mlir::OperationPass<DeviceOp>> createAIELowerCascadeFlowsPass();
 std::unique_ptr<mlir::OperationPass<DeviceOp>>
 createAIEAssignBufferDescriptorIDsPass();

include/aie/Dialect/AIE/Transforms/AIEPasses.td

@@ -344,6 +344,31 @@ def AIEObjectFifoStatefulTransform : Pass<"aie-objectFifo-stateful-transform", "
 }
 
 
+def AIEObjectFifoUnroll : Pass<"aie-objectFifo-unroll", "DeviceOp"> {
+  let summary = "Unroll scf.for loops that contain objectFifo operations";
+  let description = [{
+    Unroll each `scf.for` loop (and its ancestor loops) inside an `aie.core`
+    that contains `aie.objectfifo.acquire`/`aie.objectfifo.release` operations,
+    using the standard `scf` loop unrolling infrastructure.
+
+    Each such loop is unrolled by the least common multiple of the depths of the
+    objectFifos it accesses, so that after the objectFifo stateful transform and
+    subsequent `mem2reg`/`canonicalize` the runtime buffer/lock bookkeeping folds
+    into constant buffer references and constant lock counts (reproducing the
+    behavior of the legacy static, unrolled lowering).
+
+    This pass is intended to run before `aie-objectFifo-stateful-transform` and
+    is only used when dynamic objectFifos are disabled (`--dynamic-objFifos=false`).
+  }];
+
+  let constructor = "xilinx::AIE::createAIEObjectFifoUnrollPass()";
+  let dependentDialects = [
+    "mlir::scf::SCFDialect",
+    "xilinx::AIE::AIEDialect",
+  ];
+}
+
+
 def AIELowerCascadeFlows : Pass<"aie-lower-cascade-flows", "DeviceOp"> {
   let summary = "Lower aie.cascade_flow operations through `aie.configure_cascade` operations";
   let description = [{

lib/Dialect/AIE/Transforms/AIEObjectFifoStatefulTransform.cpp

@@ -0,0 +1,133 @@
+//===- AIEObjectFifoUnroll.cpp ----------------------------------*- C++ -*-===//
+//
+// Copyright (C) 2026 Advanced Micro Devices, Inc.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "aie/Dialect/AIE/IR/AIEDialect.h"
+#include "aie/Dialect/AIE/Transforms/AIEPasses.h"
+
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/SCF/Utils/Utils.h"
+#include "mlir/Pass/Pass.h"
+
+namespace xilinx::AIE {
+#define GEN_PASS_DEF_AIEOBJECTFIFOUNROLL
+#include "aie/Dialect/AIE/Transforms/AIEPasses.h.inc"
+} // namespace xilinx::AIE
+
+#define DEBUG_TYPE "aie-objectFifo-unroll"
+
+using namespace mlir;
+using namespace xilinx;
+using namespace xilinx::AIE;
+
+namespace {
+
+// The unroll factor for a loop is the value of the `aie.unroll_hint` attribute
+// attached by AIEObjectFifoStatefulTransform (the least common multiple of the
+// depths of the objectFifos accessed within the loop). Loops without the hint
+// carry no objectFifo access and are not unrolled (factor 1).
+static int64_t unrollFactorForLoop(scf::ForOp forOp) {
+  if (auto hint =
+          forOp->getAttrOfType<IntegerAttr>(kObjectFifoUnrollHintAttrName))
+    return hint.getInt();
+  return 1;
+}
+
+// Statically known trip count of a loop, or nullopt if it cannot be computed.
+static std::optional<int64_t> staticTripCount(scf::ForOp forOp) {
+  if (forOp.getSingleLowerBound() && forOp.getSingleUpperBound() &&
+      forOp.getSingleStep())
+    if (std::optional<llvm::APInt> tc = forOp.getStaticTripCount())
+      return tc->getSExtValue();
+  return std::nullopt;
+}
+
+// True if the loop carries an objectFifo access, i.e. it was annotated with the
+// unroll hint by the objectFifo stateful transform.
+static bool loopHasObjectFifoOp(scf::ForOp forOp) {
+  return forOp->hasAttr(kObjectFifoUnrollHintAttrName);
+}
+
+struct AIEObjectFifoUnrollPass
+    : xilinx::AIE::impl::AIEObjectFifoUnrollBase<AIEObjectFifoUnrollPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<AIEDialect>();
+    registry.insert<scf::SCFDialect>();
+  }
+
+  void runOnOperation() override {
+    DeviceOp device = getOperation();
+
+    for (auto coreOp : device.getOps<CoreOp>()) {
+      // Collect every scf.for loop that carries an objectFifo access (i.e. was
+      // annotated with the unroll hint). Ancestor loops are annotated too, so
+      // they are naturally included.
+      SmallVector<scf::ForOp> loops;
+      coreOp.walk([&](scf::ForOp forOp) {
+        if (loopHasObjectFifoOp(forOp))
+          loops.push_back(forOp);
+      });
+
+      // Operation::walk uses post-order traversal by default, so a nested loop
+      // is visited before its enclosing loop; iterating the list in order thus
+      // processes the innermost loops first. Unrolling innermost loops first
+      // avoids invalidating references to inner loops when an outer loop (which
+      // duplicates its nested loops) is unrolled.
+      for (scf::ForOp forOp : loops) {
+        int64_t unrollFactor = unrollFactorForLoop(forOp);
+        if (unrollFactor <= 1)
+          continue;
+
+        std::optional<int64_t> trip = staticTripCount(forOp);
+        // When the loop performs fewer iterations than a full rotation of the
+        // objectFifos, unroll it completely: every iteration must map to an
+        // explicit buffer/lock slot.
+        if (trip && *trip <= unrollFactor) {
+          if (failed(mlir::loopUnrollFull(forOp))) {
+            forOp.emitOpError()
+                << "failed to fully unroll objectFifo loop (trip count "
+                << *trip << ")";
+            return signalPassFailure();
+          }
+          continue;
+        }
+
+        // Otherwise unroll by the rotation period. loopUnrollByFactor peels a
+        // cleanup/epilogue loop for the remaining iterations when the trip
+        // count is not an exact multiple of the factor.
+        FailureOr<mlir::UnrolledLoopInfo> info =
+            mlir::loopUnrollByFactor(forOp, static_cast<uint64_t>(unrollFactor));
+        if (failed(info)) {
+          forOp.emitOpError()
+              << "failed to unroll objectFifo loop by factor " << unrollFactor;
+          return signalPassFailure();
+        }
+
+        // The epilogue runs the remaining (< factor) iterations. Fully unroll
+        // it as well so that each of those iterations maps to an explicit
+        // buffer/lock rotation slot. This is best-effort: an epilogue with a
+        // non-constant trip count cannot be fully unrolled and is left rolled.
+        if (info->epilogueLoopOp) {
+          scf::ForOp epilogue = *info->epilogueLoopOp;
+          if (loopHasObjectFifoOp(epilogue))
+            (void)mlir::loopUnrollFull(epilogue);
+        }
+      }
+
+      // Drop any lingering unroll hints so they do not leak into the output.
+      coreOp.walk([&](scf::ForOp forOp) {
+        forOp->removeAttr(kObjectFifoUnrollHintAttrName);
+      });
+    }
+  }
+};
+
+} // namespace
+
+std::unique_ptr<mlir::OperationPass<xilinx::AIE::DeviceOp>>
+xilinx::AIE::createAIEObjectFifoUnrollPass() {
+  return std::make_unique<AIEObjectFifoUnrollPass>();
+}

lib/Dialect/AIE/Transforms/AIEObjectFifoUnroll.cpp

@@ -24,6 +24,7 @@ add_mlir_dialect_library(
   AIEVectorTransferLowering.cpp
   AIEHoistVectorTransferPointers.cpp
   AIEObjectFifoStatefulTransform.cpp
+  AIEObjectFifoUnroll.cpp
   AIELowerCascadeFlows.cpp
   AIEGenerateColumnControlOverlay.cpp
   AIETraceToConfig.cpp

lib/Dialect/AIE/Transforms/CMakeLists.txt

@@ -0,0 +1,26 @@
+//===- uselock_dynamic_value.mlir ------------------------------*- MLIR -*-===//
+//
+// Copyright (C) 2026 Advanced Micro Devices, Inc.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+// Roundtrip test for runtime (SSA) lock values on AIE2. The lock value is a
+// general-purpose i32 register operand rather than a static attribute.
+
+// RUN: aie-opt %s | FileCheck %s
+
+module {
+  aie.device(npu1) {
+    %tile = aie.tile(1, 2)
+    %lock = aie.lock(%tile, 0) {init = 1 : i32}
+    %core = aie.core(%tile) {
+      %v = arith.constant 2 : i32
+      // CHECK: aie.use_lock(%{{.*}}, AcquireGreaterEqual, %{{.*}})
+      aie.use_lock(%lock, AcquireGreaterEqual, %v)
+      // CHECK: aie.use_lock(%{{.*}}, Release, %{{.*}})
+      aie.use_lock(%lock, Release, %v)
+      aie.end
+    }
+  }
+}