AIEObjectFifoUnroll.cpp

Go to the documentation of this file.

//===- 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/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/Transforms.h"
#include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/Dominance.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/DataLayoutInterfaces.h"
#include "mlir/Interfaces/MemorySlotInterfaces.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Transforms/Mem2Reg.h"
#include "mlir/Transforms/Passes.h"
 
#include "llvm/ADT/DenseSet.h"
 
#include <limits>
 
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;
 
bool xilinx::AIE::isForeverLoop(scf::WhileOp whileOp) {
  scf::ConditionOp condition = whileOp.getConditionOp();
  return matchPattern(condition.getCondition(), m_One()) &&
         llvm::equal(condition.getArgs(), whileOp.getBeforeArguments()) &&
         llvm::all_of(whileOp.getBeforeBody()->without_terminator(),
                      [](Operation &op) { return isPure(&op); });
}
 
bool xilinx::AIE::isForeverForLoop(scf::ForOp forOp) {
  return forOp.getInductionVar().use_empty() &&
         isConstantIntValue(forOp.getLowerBound(), 0) &&
         isConstantIntValue(forOp.getStep(), 1) &&
         isConstantIntValue(forOp.getUpperBound(),
                            std::numeric_limits<int64_t>::max());
}
 
scf::WhileOp xilinx::AIE::replaceWithForeverLoop(scf::ForOp forOp) {
  OpBuilder builder(forOp);
  auto whileOp = scf::WhileOp::create(
      builder, forOp.getLoc(), forOp.getResultTypes(), forOp.getInitArgs(),
      [](OpBuilder &b, Location loc, ValueRange args) {
        Value always = arith::ConstantIntOp::create(b, loc, 1, 1);
        scf::ConditionOp::create(b, loc, always, args);
      },
      [](OpBuilder &, Location, ValueRange) {});
  Block *forBody = forOp.getBody();
  Block *after = whileOp.getAfterBody();
  for (auto [from, to] :
       llvm::zip_equal(forOp.getRegionIterArgs(), after->getArguments())) {
    from.replaceAllUsesWith(to);
  }
  after->getOperations().splice(after->end(), forBody->getOperations());
  forOp.replaceAllUsesWith(whileOp.getResults());
  forOp.erase();
  return whileOp;
}
 
bool xilinx::AIE::isObjectFifoLoop(Operation *op) {
  if (auto whileOp = dyn_cast<scf::WhileOp>(op)) {
    return isForeverLoop(whileOp);
  }
  return isa<scf::ForOp>(op);
}
 
Operation *xilinx::AIE::getObjectFifoLoopParent(Operation *op) {
  for (Operation *parent = op->getParentOp(); parent;
       parent = parent->getParentOp()) {
    if (isObjectFifoLoop(parent)) {
      return parent;
    }
  }
  return nullptr;
}
 
namespace {
 
// Statically known trip count of a loop, or nullopt if it cannot be computed.
static std::optional<int64_t> getStaticTripCount(scf::ForOp forOp) {
  if (forOp.getSingleLowerBound() && forOp.getSingleUpperBound() &&
      forOp.getSingleStep()) {
    if (std::optional<llvm::APInt> tc = forOp.getStaticTripCount()) {
      return tc->getSExtValue();
    }
  }
  return std::nullopt;
}
 
// Unroll a forever loop by `factor`: its body runs `factor` times per trip
// around the loop, each copy taking the values the previous one yields. The
// loop never exits, so there is no remainder to handle.
static void unrollForeverLoop(scf::WhileOp whileOp, int64_t factor) {
  Block *body = whileOp.getAfterBody();
  auto yield = cast<scf::YieldOp>(body->getTerminator());
  SmallVector<Operation *> bodyOps;
  for (Operation &op : body->without_terminator()) {
    bodyOps.push_back(&op);
  }
  OpBuilder builder(yield);
  SmallVector<Value> yielded(yield.getOperands());
  SmallVector<Value> carried = yielded;
  for (int64_t i = 1; i < factor; ++i) {
    IRMapping mapping;
    mapping.map(body->getArguments(), carried);
    for (Operation *op : bodyOps) {
      builder.clone(*op, mapping);
    }
    for (auto [next, value] : llvm::zip_equal(carried, yielded)) {
      next = mapping.lookupOrDefault(value);
    }
  }
  yield->setOperands(carried);
}
 
// Remove redundant AIE1 binary-lock acquires.
//
// The dynamic lowering acquires the whole sliding window ([counter, counter +
// acqNumber)) on every loop iteration; after unrolling and constant folding
// that expands to an explicit `Acquire` of each element's lock.
// This drops the needless re-acquire of locks already held.
//
// The analysis is block-local and conservatively assumes nothing is held on
// block entry. That keeps the very first acquire of each unrolled loop body
// (which legitimately re-establishes the window across the back-edge) while
// dropping the intra-body re-acquires of locks whose element is still held.
static bool removeRedundantBinaryAcquires(DeviceOp device) {
  bool erased = false;
  device.walk([&](Block *block) {
    llvm::DenseSet<Value> held;
    SmallVector<UseLockOp> toErase;
    for (Operation &op : *block) {
      auto useLock = dyn_cast<UseLockOp>(&op);
      if (!useLock) {
        continue;
      }
      Value lock = useLock.getLock();
      if (useLock.acquire()) {
        // A binary `Acquire` of a lock already in the held set is redundant.
        if (!held.insert(lock).second) {
          toErase.push_back(useLock);
        }
      } else if (useLock.release()) {
        held.erase(lock);
      }
    }
    for (UseLockOp op : toErase) {
      op.erase();
    }
    erased |= !toErase.empty();
  });
  return erased;
}
 
// Marks a loop already offered to the peeling trial below, so that neither the
// peeled loop nor the reverted copy is considered again.
static constexpr llvm::StringLiteral kPeelTriedAttrName = "aie.peel_tried";
 
// Number of `use_lock` ops in `op` whose lock value is not a compile-time
// constant. UseLockOp::getConstantValue() reports an error when it fails, so
// this probes the value directly.
static int64_t countRuntimeLockValues(Operation *op) {
  int64_t count = 0;
  op->walk([&](UseLockOp useLockOp) {
    if (!getConstantIntValue(useLockOp.getValue())) {
      count++;
    }
  });
  return count;
}
 
// Number of memrefs in `op` that are picked at run time: block arguments and
// the results of region ops and selects.
static int64_t countRuntimeMemrefs(Operation *op) {
  int64_t count = 0;
  auto isMemref = llvm::IsaPred<BaseMemRefType>;
  op->walk([&](Block *block) {
    count += llvm::count_if(block->getArgumentTypes(), isMemref);
  });
  op->walk([&](Operation *nested) {
    if (nested != op &&
        (nested->getNumRegions() > 0 || isa<arith::SelectOp>(nested))) {
      count += llvm::count_if(nested->getResultTypes(), isMemref);
    }
  });
  return count;
}
 
// Canonicalize, SCCP, canonicalize. The fold needs the folding patterns only:
// region simplification walks every region of the device, runtime sequences
// included, and the canonicalizes after this pass still run it.
static OpPassManager buildFoldPipeline() {
  GreedyRewriteConfig config;
  config.setRegionSimplificationLevel(GreedySimplifyRegionLevel::Disabled);
  OpPassManager pipeline(DeviceOp::getOperationName());
  pipeline.addPass(mlir::createCanonicalizerPass(config));
  pipeline.addPass(mlir::createSCCPPass());
  pipeline.addPass(mlir::createCanonicalizerPass(config));
  return pipeline;
}
 
struct AIEObjectFifoUnrollPass
    : xilinx::AIE::impl::AIEObjectFifoUnrollBase<AIEObjectFifoUnrollPass> {
  using Base::Base;
 
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<AIEDialect>();
    registry.insert<scf::SCFDialect>();
  }
 
  /// Peel the first iteration off each loop whose lock values are still
  /// computed at run time, keeping a peel only where it makes them constant.
  ///
  /// Whether a peel pays off depends on foldings that cannot be predicted from
  /// the objectFifo access pattern alone (a conditional acquire may or may not
  /// collapse, for instance), so each candidate is peeled for real and the
  /// result measured. While the fold pipeline runs, the unpeeled core is held
  /// aside as a detached clone, so exactly one version of the loop is ever
  /// live; that clone is swapped back in when the peel bought nothing, and
  /// destroyed otherwise.
  ///
  /// The measurement is a per-core count, so attributing a change to a peel
  /// needs one peel per core per fold; the loop below drains the candidates one
  /// round at a time, peeling every core in parallel so a round costs a single
  /// run of the fold pipeline. Candidates are marked before they are tried and
  /// never unmarked while draining, so the number of rounds is bounded by the
  /// most objectFifo loops in any one core.
  LogicalResult peelObjectFifoLoops(DeviceOp device) {
    OpPassManager foldPipeline = buildFoldPipeline();
 
    IRRewriter rewriter(device.getContext());
    struct Trial {
      CoreOp core;
      Operation *unpeeled;
      int64_t runtimeLocksBefore;
    };
 
    while (true) {
      SmallVector<Trial> trials;
      for (CoreOp core : SmallVector<CoreOp>(device.getOps<CoreOp>())) {
        scf::ForOp target;
        // Post-order, so the innermost untried loop is taken first.
        core.walk([&](scf::ForOp forOp) {
          if (forOp->hasAttr(kPeelTriedAttrName) ||
              countRuntimeLockValues(forOp) == 0) {
            return WalkResult::advance();
          }
          // A single-iteration loop has nothing to peel off.
          std::optional<int64_t> trip = getStaticTripCount(forOp);
          if (trip && *trip <= 1) {
            return WalkResult::advance();
          }
          target = forOp;
          return WalkResult::interrupt();
        });
        if (!target) {
          continue;
        }
        // Marked before the core is cloned, so neither the peeled loop nor the
        // fallback copy is offered as a candidate again.
        target->setAttr(kPeelTriedAttrName, rewriter.getUnitAttr());
 
        Trial trial{core, core->clone(), countRuntimeLockValues(core)};
        scf::ForOp firstIteration;
        if (failed(scf::peelForLoopFirstIteration(rewriter, target,
                                                  firstIteration))) {
          trial.unpeeled->erase();
          continue;
        }
        trials.push_back(trial);
      }
      if (trials.empty()) {
        break;
      }
 
      if (failed(runPipeline(foldPipeline, device))) {
        for (Trial &trial : trials) {
          trial.unpeeled->erase();
        }
        return failure();
      }
 
      for (Trial &trial : trials) {
        if (countRuntimeLockValues(trial.core) < trial.runtimeLocksBefore) {
          trial.unpeeled->erase();
          continue;
        }
        rewriter.setInsertionPoint(trial.core);
        Operation *restored = rewriter.insert(trial.unpeeled);
        rewriter.replaceOp(trial.core, restored->getResults());
      }
    }
 
    device.walk(
        [&](scf::ForOp forOp) { forOp->removeAttr(kPeelTriedAttrName); });
    return success();
  }
 
  /// Run `foldPipeline`, first promoting the object slots collected above to
  /// SSA on every core where that folds each object to its buffer; elsewhere
  /// they are left in memory.
  ///
  /// A promoted object still picked at run time is a memref descriptor carried
  /// through control flow, which LLVM's alias analysis cannot tell apart from
  /// the others; left in memory, SROA reduces it to a pointer. As with peeling,
  /// each core is promoted for real against a detached clone. Unless the
  /// promoted core picks no memref at run time, the clone is swapped in and
  /// folded too, and the promoted core returns only if it picks no more: with
  /// both sides folded, unrelated folding cannot tip the comparison.
  LogicalResult foldPromotingObjectSlots(DeviceOp device,
                                         ArrayRef<memref::AllocaOp> slots,
                                         OpPassManager &foldPipeline) {
    if (slots.empty()) {
      return runPipeline(foldPipeline, device);
    }
    struct Trial {
      CoreOp core;
      Operation *unpromoted;
    };
    SmallVector<Trial> trials;
    llvm::SmallPtrSet<Operation *, 8> seen;
    SmallVector<PromotableAllocationOpInterface> allocators;
    for (memref::AllocaOp slot : slots) {
      auto core = slot->getParentOfType<CoreOp>();
      if (seen.insert(core).second) {
        trials.push_back({core, core->clone()});
      }
      allocators.push_back(
          cast<PromotableAllocationOpInterface>(slot.getOperation()));
    }
    DataLayout dataLayout = DataLayout::closest(device);
    DominanceInfo dominance(device);
    OpBuilder builder(device.getContext());
    (void)tryToPromoteMemorySlots(allocators, builder, dataLayout, dominance);
 
    if (failed(runPipeline(foldPipeline, device))) {
      for (Trial &trial : trials) {
        trial.unpromoted->erase();
      }
      return failure();
    }
 
    IRRewriter rewriter(device.getContext());
    // Puts `in` in place of `out` and detaches `out`.
    auto swap = [&](Operation *out, Operation *in) {
      rewriter.setInsertionPoint(out);
      rewriter.insert(in);
      rewriter.replaceAllOpUsesWith(out, in->getResults());
      out->remove();
    };
    struct Contest {
      Operation *promoted;
      Operation *unpromoted;
      int64_t promotedRuntimeMemrefs;
    };
    SmallVector<Contest> contests;
    bool restoredAny = false;
    for (Trial &trial : trials) {
      // mem2reg may also have left a slot in place.
      WalkResult leftover = trial.core.walk([&](memref::AllocaOp allocaOp) {
        return allocaOp->hasAttr(kObjectFifoObjectSlotAttrName)
                   ? WalkResult::interrupt()
                   : WalkResult::advance();
      });
      int64_t runtimeMemrefs = countRuntimeMemrefs(trial.core);
      if (!leftover.wasInterrupted() && runtimeMemrefs == 0) {
        trial.unpromoted->erase();
        continue;
      }
      swap(trial.core, trial.unpromoted);
      trial.unpromoted->walk([&](memref::AllocaOp allocaOp) {
        allocaOp->removeAttr(kObjectFifoObjectSlotAttrName);
      });
      if (leftover.wasInterrupted()) {
        trial.core->erase();
      } else {
        contests.push_back({trial.core, trial.unpromoted, runtimeMemrefs});
      }
      restoredAny = true;
    }
    if (!restoredAny) {
      return success();
    }
 
    // The restored cores have not been folded yet.
    if (failed(runPipeline(foldPipeline, device))) {
      for (Contest &contest : contests) {
        contest.promoted->erase();
      }
      return failure();
    }
    for (Contest &contest : contests) {
      if (contest.promotedRuntimeMemrefs <=
          countRuntimeMemrefs(contest.unpromoted)) {
        swap(contest.unpromoted, contest.promoted);
        contest.unpromoted->erase();
      } else {
        contest.promoted->erase();
      }
    }
    return success();
  }
 
  void runOnOperation() override {
    DeviceOp device = getOperation();
 
    bool peelFirstIteration = false;
    if (clPeelFirstIteration == "auto") {
      peelFirstIteration = !clDefaultDynamic;
    } else if (clPeelFirstIteration == "true") {
      peelFirstIteration = true;
    } else if (clPeelFirstIteration != "false") {
      device.emitOpError("invalid peel-first-iteration value '")
          << clPeelFirstIteration << R"('; expected "true", "false" or "auto")";
      return signalPassFailure();
    }
 
    SmallVector<memref::AllocaOp> objectSlots;
    for (auto coreOp : device.getOps<CoreOp>()) {
      // `default-dynamic` picks the lowering for cores that do not pin their
      // own choice; a core's `dynamic_objfifo_lowering` attribute overrides it
      // in either direction. Dynamic cores keep their loops rolled (drop the
      // hints, skip unrolling) and keep in memory the objects they reach from
      // anywhere but straight-line code; their runtime bookkeeping is tidied by
      // the shared fold/cleanup pipeline below.
      if (coreOp.getDynamicObjfifoLowering().value_or(clDefaultDynamic)) {
        coreOp.walk([&](Operation *op) {
          op->removeAttr(kObjectFifoUnrollHintAttrName);
        });
        bool straightLine = coreOp.getBody().hasOneBlock();
        coreOp.walk([&](memref::AllocaOp allocaOp) {
          if (!allocaOp->hasAttr(kObjectFifoObjectSlotAttrName)) {
            return;
          }
          if (straightLine &&
              llvm::all_of(allocaOp->getUsers(), [&](Operation *user) {
                return user->getParentOp() == coreOp.getOperation();
              })) {
            objectSlots.push_back(allocaOp);
          } else {
            allocaOp->removeAttr(kObjectFifoObjectSlotAttrName);
          }
        });
        continue;
      }
      coreOp.walk([&](memref::AllocaOp allocaOp) {
        if (allocaOp->hasAttr(kObjectFifoObjectSlotAttrName)) {
          objectSlots.push_back(allocaOp);
        }
      });
      SmallVector<Operation *> loops;
      coreOp.walk([&](Operation *op) {
        if (isObjectFifoLoop(op)) {
          loops.push_back(op);
        }
      });
 
      // 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 (Operation *loop : loops) {
        // Loops without the objectFifo unroll hint (the LCM of the accessed
        // fifo depths) carry no objectFifo access; factor 1 leaves them rolled.
        auto hint =
            loop->getAttrOfType<IntegerAttr>(kObjectFifoUnrollHintAttrName);
        int64_t unrollFactor = hint ? hint.getInt() : 1;
        if (unrollFactor <= 1) {
          continue;
        }
 
        // A forever loop has no trip count to cover, so repeating its body
        // once per rotation is enough.
        auto forOp = dyn_cast<scf::ForOp>(loop);
        if (!forOp) {
          unrollForeverLoop(cast<scf::WhileOp>(loop), unrollFactor);
          continue;
        }
        if (isForeverForLoop(forOp)) {
          unrollForeverLoop(replaceWithForeverLoop(forOp), unrollFactor);
          continue;
        }
 
        std::optional<int64_t> trip = getStaticTripCount(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.
        std::optional<scf::ForOp> epilogue =
            info->epilogueLoopOp; // NOLINT(bugprone-unchecked-optional-access)
        if (epilogue) {
          (void)mlir::loopUnrollFull(*epilogue);
        }
      }
 
      // Drop any lingering unroll hints so they do not leak into the output.
      coreOp.walk([&](Operation *op) {
        op->removeAttr(kObjectFifoUnrollHintAttrName);
      });
    }
 
    // --aie-objectfifo-lower-cores promotes the lock bookkeeping counters to
    // loop-carried SSA values, and the objects are promoted here where they
    // fold. Once the loops have been unrolled by their rotation period, the
    // objects and counters become loop-invariant, so every object and lock
    // value collapses to a constant. Run that fold here as a scoped
    // sub-pipeline: canonicalize exposes the constants, SCCP propagates them
    // across any remainder loop that survives a partial unroll, and a final
    // canonicalize deletes the now-dead counter arithmetic and iter_args.
    OpPassManager foldPipeline = buildFoldPipeline();
    if (failed(foldPromotingObjectSlots(device, objectSlots, foldPipeline))) {
      return signalPassFailure();
    }
 
    if (peelFirstIteration && failed(peelObjectFifoLoops(device))) {
      return signalPassFailure();
    }
 
    // With the window acquires now folded to concrete per-lock `Acquire` ops,
    // drop the ones that re-acquire an already-held AIE1 binary lock. Only a
    // drop leaves dead constants for a final canonicalize; the fold left the
    // device canonical otherwise.
    if (!removeRedundantBinaryAcquires(device)) {
      return;
    }
    OpPassManager cleanupPipeline(DeviceOp::getOperationName());
    cleanupPipeline.addPass(mlir::createCanonicalizerPass());
    if (failed(runPipeline(cleanupPipeline, device))) {
      return signalPassFailure();
    }
  }
};
 
} // namespace
 
std::unique_ptr<mlir::OperationPass<xilinx::AIE::DeviceOp>>
xilinx::AIE::createAIEObjectFifoUnrollPass() {
  return std::make_unique<AIEObjectFifoUnrollPass>();
}
 
std::unique_ptr<OperationPass<DeviceOp>>
xilinx::AIE::createAIEObjectFifoUnrollPass(
    const AIEObjectFifoUnrollOptions &options) {
  return std::make_unique<AIEObjectFifoUnrollPass>(options);
}