AIEObjectFifoStatefulTransform.cpp

//===- AIEObjectFifoStatefulTransform.cpp ----------------------*- MLIR -*-===//
//
// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
// (c) Copyright 2021 Xilinx Inc.
//
// Date: October 18th 2021
//
//===----------------------------------------------------------------------===//

#include "aie/Dialect/AIE/IR/AIEDialect.h"
#include "aie/Dialect/AIE/Transforms/AIEPasses.h"

#include "mlir/Analysis/TopologicalSortUtils.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/Utils/Utils.h"
#include "mlir/IR/Attributes.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/DialectConversion.h"

#include "mlir/IR/Operation.h"
#include "mlir/Interfaces/DataLayoutInterfaces.h"

#include <numeric>
#include <set>

#include <iostream>

namespace xilinx::AIE {
#define GEN_PASS_DEF_AIEOBJECTFIFOSTATEFULTRANSFORM
#include "aie/Dialect/AIE/Transforms/AIEPasses.h.inc"
} // namespace xilinx::AIE

using namespace mlir;
using namespace xilinx;
using namespace xilinx::AIE;

#define DEBUG_TYPE "aie-objectFifo-stateful-transform"

#define LOOP_VAR_DEPENDENCY (-2)

//===----------------------------------------------------------------------===//
// Lock Analysis
//===----------------------------------------------------------------------===//
class LockAnalysis {
  DenseMap<std::pair<Value, int>, int> locksPerTile;

public:
  LockAnalysis(DeviceOp &device) {
    // go over the locks created for each tile and update the index in
    // locksPerTile
    device.walk([&](LockOp lockOp) {
      auto tile = lockOp.getTile();
      auto lockID = lockOp.getLockIDValue();
      locksPerTile[{tile, lockID}] = 1;
    });
  }

  /// Given a tile, returns next usable lockID for that tile.
  int getLockID(TileOp &tileOp) {
    const auto &targetModel = getTargetModel(tileOp);
    for (unsigned i = 0;
         i < targetModel.getNumLocks(tileOp.getCol(), tileOp.getRow()); i++)
      if (int usageCnt = locksPerTile[{tileOp, i}]; usageCnt == 0) {
        locksPerTile[{tileOp, i}] = 1;
        return i;
      }
    return -1;
  }
};

//===----------------------------------------------------------------------===//
// DMA Channel Analysis
//===----------------------------------------------------------------------===//
class DMAChannelAnalysis {
  DenseMap<std::tuple<Value, DMAChannelDir, int>, int> channelsPerTile;
  DenseMap<std::tuple<Value, DMAChannelDir, int>, int> aieStreamsPerTile;

public:
  DMAChannelAnalysis(DeviceOp &device) {
    // go over the channels used for each tile and update channel map
    for (auto memOp : device.getOps<MemOp>()) {
      Region &r = memOp.getBody();
      for (auto &bl : r.getBlocks()) {
        for (auto op : bl.getOps<DMAStartOp>()) {
          channelsPerTile[{memOp.getTile(), op.getChannelDir(),
                           op.getChannelIndex()}] = 1;
        }
      }
    }
    for (auto memOp : device.getOps<MemTileDMAOp>()) {
      Region &r = memOp.getBody();
      for (auto &bl : r.getBlocks()) {
        for (auto op : bl.getOps<DMAStartOp>()) {
          channelsPerTile[{memOp.getTile(), op.getChannelDir(),
                           op.getChannelIndex()}] = 1;
        }
      }
    }
    for (auto memOp : device.getOps<ShimDMAOp>()) {
      Region &r = memOp.getBody();
      for (auto &bl : r.getBlocks()) {
        for (auto op : bl.getOps<DMAStartOp>()) {
          channelsPerTile[{memOp.getTile(), op.getChannelDir(),
                           op.getChannelIndex()}] = 1;
        }
      }
    }
    for (auto flowOp : device.getOps<FlowOp>()) {
      if (flowOp.getSourceBundle() == WireBundle::Core)
        aieStreamsPerTile[{flowOp.getSource(), DMAChannelDir::MM2S,
                           flowOp.getSourceChannel()}] = 1;
      if (flowOp.getDestBundle() == WireBundle::Core)
        aieStreamsPerTile[{flowOp.getDest(), DMAChannelDir::S2MM,
                           flowOp.getDestChannel()}] = 1;
    }
  }

  /// Given a tile and DMAChannelDir, returns next usable channel index for
  /// that tile.
  int getDMAChannelIndex(TileOp tileOp, DMAChannelDir dir,
                         bool requiresAdjacentTileAccessChannels) {
    int maxChannelNum = 0;
    if (dir == DMAChannelDir::MM2S)
      maxChannelNum = tileOp.getNumSourceConnections(WireBundle::DMA);
    else
      maxChannelNum = tileOp.getNumDestConnections(WireBundle::DMA);

    const auto &targetModel = getTargetModel(tileOp);
    int maxChannelNumForAdjacentTile =
        targetModel.getMaxChannelNumForAdjacentMemTile(tileOp.getCol(),
                                                       tileOp.getRow());

    // if requires adjacent tile access channels, only allocate on channel 0-3,
    // and if cannot, return 0
    if (requiresAdjacentTileAccessChannels) {
      maxChannelNum = std::min(maxChannelNum, maxChannelNumForAdjacentTile);
    }

    for (int i = 0; i < maxChannelNum; i++) {
      if (int usageCnt = channelsPerTile[{tileOp.getResult(), dir, i}];
          usageCnt == 0) {
        channelsPerTile[{tileOp.getResult(), dir, i}] = 1;
        return i;
      }
    }
    return -1;
  }

  /// Given a tile and DMAChannel, adds entry to aieStreamsPerTile or
  /// throws an error if the stream is already used.
  void checkAIEStreamIndex(TileOp tileOp, DMAChannel chan) {
    if (aieStreamsPerTile.find({tileOp.getResult(), chan.direction,
                                chan.channel}) == aieStreamsPerTile.end()) {
      aieStreamsPerTile[{tileOp.getResult(), chan.direction, chan.channel}] = 1;
    } else {
      if (chan.direction == DMAChannelDir::MM2S)
        tileOp.emitOpError("number of output Core channels exceeded!");
      else
        tileOp.emitOpError("number of input Core channels exceeded!");
    }
  }
};

//===----------------------------------------------------------------------===//
// Create objectFifos Pass
//===----------------------------------------------------------------------===//

/// Struct to hold per-device state for the objectFifo transformation.
/// This is passed to helper functions to avoid member variable pollution
/// between different device operations.
struct ObjectFifoState {
  DenseMap<ObjectFifoCreateOp, std::vector<BufferOp>>
      buffersPerFifo; // maps each objFifo to its corresponding buffer
  DenseMap<ObjectFifoCreateOp, std::vector<ExternalBufferOp>>
      externalBuffersPerFifo; // maps each objFifo to its corresponding
                              // external buffers
  DenseMap<ObjectFifoCreateOp, std::vector<LockOp>>
      locksPerFifo; // maps each objFifo to its corresponding locks
  std::vector<std::pair<ObjectFifoCreateOp, std::vector<ObjectFifoCreateOp>>>
      splitFifos; // maps each objFifo between non-adjacent tiles to its
                  // corresponding consumer objectFifos
  DenseMap<ObjectFifoLinkOp, ObjectFifoCreateOp>
      objFifoLinks; // maps each ObjectFifoLinkOp to objFifo whose elements
                    // have been created and should be used
  std::vector<ObjectFifoCreateOp>
      splitBecauseLink; // objfifos which have been split because they are
                        // part of a Link, not because they didn't have a shared
                        // memory module
};

struct AIEObjectFifoStatefulTransformPass
    : xilinx::AIE::impl::AIEObjectFifoStatefulTransformBase<
          AIEObjectFifoStatefulTransformPass> {

  /// Function that returns true if two tiles in the AIE array share a memory
  /// module. share_direction is equal to:
  ///   * 2 if the memory modules on both tiles can be shared,
  ///   * -1 if the shared memory module is that of the first input tile,
  ///   * 1 if it is that of the second input tile,
  ///   * 0 is no memory module is shared.
  bool isSharedMemory(TileOp a, TileOp b, int *share_direction) {
    const auto &targetModel = getTargetModel(a.getOperation());

    if ((a.isShimTile() && !b.isShimTile()) ||
        (!a.isShimTile() && b.isShimTile())) {
      *share_direction = 0;
      return false;
    }
    if ((targetModel.isMemTile(a.getCol(), a.getRow()) &&
         !targetModel.isMemTile(b.getCol(), b.getRow())) ||
        (!targetModel.isMemTile(a.getCol(), a.getRow()) &&
         targetModel.isMemTile(b.getCol(), b.getRow()))) {
      *share_direction = 0;
      return false;
    }
    bool rightShared = targetModel.isLegalMemAffinity(
        a.colIndex(), a.rowIndex(), b.colIndex(), b.rowIndex());

    bool leftShared = targetModel.isLegalMemAffinity(
        b.colIndex(), b.rowIndex(), a.colIndex(), a.rowIndex());

    if (leftShared && rightShared)
      *share_direction = 2;
    else if (leftShared)
      *share_direction = -1;
    else if (rightShared)
      *share_direction = 1;
    else
      *share_direction = 0;

    return leftShared || rightShared;
  }

  /// Function to retrieve ObjectFifoAllocateOp of ObjectFifoCreateOp,
  /// if it exists.
  std::optional<ObjectFifoAllocateOp>
  getOptionalAllocateOp(ObjectFifoCreateOp op) {
    ObjectFifoAllocateOp allocOp;
    auto device = op->getParentOfType<DeviceOp>();
    bool foundAlloc = false;
    for (ObjectFifoAllocateOp alloc : device.getOps<ObjectFifoAllocateOp>()) {
      if (alloc.getObjectFifo() == op) {
        if (foundAlloc)
          op.emitOpError("has more than one allocate operation");
        allocOp = alloc;
        foundAlloc = true;
      }
    }
    if (foundAlloc)
      return {allocOp};
    return {};
  }

  // Return true if the objectFifo created by createOp requires a DMA to be set
  // up. This is the case if the tiles are not adjacent (no shared memory), if
  // the objectFifo broadcasts to multiple tiles, if one of the consumers or
  // the producer wants to use the multi-dimensional address generation
  // features of the DMA, if the objectFifo is part of a LinkOp, or if the
  // via_DMA or repeatCount attributes of the objectFifo are set.
  bool requiresDMAs(ObjectFifoCreateOp createOp, int &share_direction,
                    ObjectFifoState &state) {
    bool hasSharedMemory = false;
    bool atLeastOneConsumerWantsTransform = false;
    bool isUsedInLinkOp = false;

    if (createOp.getVia_DMA())
      return true;

    if (createOp.getRepeatCount().has_value())
      return true;

    if (createOp.getAieStream())
      return true;

    if (createOp.getConsumerTiles().size() == 1 &&
        createOp.getDimensionsToStream().empty()) {

      // Test for shared memory
      for (auto consumerTile : createOp.getConsumerTiles()) {
        if (auto consumerTileOp =
                dyn_cast<TileOp>(consumerTile.getDefiningOp())) {
          if (std::count(state.splitBecauseLink.begin(),
                         state.splitBecauseLink.end(), createOp))
            hasSharedMemory =
                isSharedMemory(createOp.getProducerTileOp(),
                               createOp.getProducerTileOp(), &share_direction);
          else
            hasSharedMemory = isSharedMemory(createOp.getProducerTileOp(),
                                             consumerTileOp, &share_direction);
        }
      }
    }

    // Only test for use of data layout transformations if we are in the shared
    // memory case; otherwise, we will return `true` in any case.
    if (hasSharedMemory) {
      // Even if just one of the consumers in the list of consumers wants to
      // perform a memory transform, we need to use DMAs.
      for (BDDimLayoutArrayAttr dims :
           createOp.getDimensionsFromStreamPerConsumer())
        if (!dims.empty()) {
          atLeastOneConsumerWantsTransform = true;
          break;
        }
    }

    // Check if the objectfifo operation can use shared memory for linking. If
    // the link operation is a distribute or a join operation, or if the link
    // has different memref types, DMAs are required even if shared memory is
    // available and the objectfifo should be split. Otherwise also check if the
    // via_shared_memory attribute of the objectfifo operation is set and try to
    // apply it.
    if (hasSharedMemory) {
      if (auto linkOp = getOptionalLinkOp(createOp)) {
        isUsedInLinkOp = true;
        if (!linkOp->isDistribute() && !linkOp->isJoin()) {
          auto fifoInType = llvm::cast<AIEObjectFifoType>(
              linkOp->getInputObjectFifos()[0].getElemType());
          auto producerType =
              llvm::cast<MemRefType>(fifoInType.getElementType());
          auto fifoOutType = llvm::cast<AIEObjectFifoType>(
              linkOp->getOutputObjectFifos()[0].getElemType());
          auto consumerType =
              llvm::cast<MemRefType>(fifoOutType.getElementType());
          if (consumerType != producerType) {
            // TODO: Support for different memref types through shared
            // memory without DMAs
            state.splitBecauseLink.push_back(createOp);
          }
          std::optional<ObjectFifoAllocateOp> opAlloc =
              getOptionalAllocateOp(createOp);
          if (opAlloc.has_value()) {
            TileOp delegate = opAlloc->getDelegateTileOp();
            int prodShareDir;
            int consShareDir;
            auto consumerTileOp = dyn_cast<TileOp>(
                createOp.getConsumerTiles()[0].getDefiningOp());
            isSharedMemory(delegate, createOp.getProducerTileOp(),
                           &prodShareDir);
            isSharedMemory(delegate, consumerTileOp, &consShareDir);
            if ((prodShareDir == -1 || prodShareDir == 2) &&
                (consShareDir == -1 || consShareDir == 2))
              isUsedInLinkOp = false;
            else
              state.splitBecauseLink.push_back(createOp);
          }
        } else {
          state.splitBecauseLink.push_back(createOp);
        }
      }
    }

    return !hasSharedMemory || atLeastOneConsumerWantsTransform ||
           isUsedInLinkOp;
  }

  /// Function to retrieve ObjectFifoLinkOp of ObjectFifoCreateOp,
  /// if it belongs to one.
  std::optional<ObjectFifoLinkOp> getOptionalLinkOp(ObjectFifoCreateOp op) {
    auto device = op->getParentOfType<DeviceOp>();
    for (ObjectFifoLinkOp linkOp : device.getOps<ObjectFifoLinkOp>()) {
      for (ObjectFifoCreateOp in : linkOp.getInputObjectFifos())
        if (in == op)
          return {linkOp};
      for (ObjectFifoCreateOp out : linkOp.getOutputObjectFifos())
        if (out == op)
          return {linkOp};
    }
    return {};
  }

  ObjectFifoCreateOp
  createObjectFifo(OpBuilder &builder, AIEObjectFifoType datatype,
                   std::string name, Value prodTile, Value consTile,
                   Attribute depth, BDDimLayoutArrayAttr dimensionsToStream,
                   BDDimLayoutArrayArrayAttr dimensionsFromStreamPerConsumer) {
    auto ofName = builder.getStringAttr(name);
    auto fifo = ObjectFifoCreateOp::create(
        builder, builder.getUnknownLoc(), ofName, prodTile, consTile, depth,
        datatype, dimensionsToStream, dimensionsFromStreamPerConsumer);
    return fifo;
  }

  /// Function used to create objectFifo locks based on target architecture.
  /// Called by createObjectFifoElements().
  std::vector<LockOp>
  createObjectFifoLocks(OpBuilder &builder, LockAnalysis &lockAnalysis,
                        ObjectFifoCreateOp op, int numElem,
                        int joinDistribFactor, TileOp creation_tile,
                        int repeatCount, ObjectFifoState &state) {
    std::vector<LockOp> locks;
    if (op.getDisableSynchronization())
      return locks;
    auto dev = op->getParentOfType<DeviceOp>();
    auto &target = dev.getTargetModel();
    // if shimTile external buffers are collected from input code
    // create as many locks as there are external buffers
    if (creation_tile.isShimTile()) {
      numElem = 0;
      if (!state.externalBuffersPerFifo[op].empty())
        numElem = state.externalBuffersPerFifo[op].size();
    }
    if (target.getTargetArch() == AIEArch::AIE1) {
      for (int i = 0; i < numElem; i++) {
        // create corresponding aie1 locks
        int initValue = op.getInitValues().has_value() ? 1 : 0;
        int lockID = lockAnalysis.getLockID(creation_tile);
        assert(lockID >= 0 && "No more locks to allocate!");
        auto lock = LockOp::create(builder, builder.getUnknownLoc(),
                                   creation_tile, lockID, initValue);
        lock.getOperation()->setAttr(SymbolTable::getSymbolAttrName(),
                                     builder.getStringAttr(op.name().str() +
                                                           "_lock_" +
                                                           std::to_string(i)));
        locks.push_back(lock);
      }
    } else {
      // create corresponding aie2 locks
      for (int i = 0; i < joinDistribFactor; i++) {
        auto initValues = op.getInitValues().has_value()
                              ? op.getInitValues().value().size()
                              : 0;
        int prodLockID = lockAnalysis.getLockID(creation_tile);
        assert(prodLockID >= 0 && "No more locks to allocate!");
        int prodLockValue = (numElem - initValues) * repeatCount;
        auto prodLock =
            LockOp::create(builder, builder.getUnknownLoc(), creation_tile,
                           prodLockID, prodLockValue);
        prodLock.getOperation()->setAttr(
            SymbolTable::getSymbolAttrName(),
            builder.getStringAttr(op.name().str() + "_prod_lock_" +
                                  std::to_string(i)));
        locks.push_back(prodLock);

        int consLockID = lockAnalysis.getLockID(creation_tile);
        assert(consLockID >= 0 && "No more locks to allocate!");
        int consLockValue = initValues * repeatCount;
        auto consLock =
            LockOp::create(builder, builder.getUnknownLoc(), creation_tile,
                           consLockID, consLockValue);
        consLock.getOperation()->setAttr(
            SymbolTable::getSymbolAttrName(),
            builder.getStringAttr(op.name().str() + "_cons_lock_" +
                                  std::to_string(i)));
        locks.push_back(consLock);
      }
    }
    return locks;
  }

  /// Function to calculate total memory usage on a specific tile
  /// based on all buffers allocated to that tile from buffersPerFifo map
  int calculateCurrentUsedMemory(
      TileOp targetTile,
      DenseMap<ObjectFifoCreateOp, std::vector<BufferOp>> &buffersPerFifo,
      std::vector<BufferOp> &buffers) {
    int totalUsedMemory = 0;

    // Iterate through all ObjectFifos and their buffers
    for (auto &[fifoOp, bufferList] : buffersPerFifo) {
      for (auto &buffer : bufferList) {
        // Check if this buffer is allocated on the target tile
        if (buffer.getTile() == targetTile.getResult()) {
          auto bufferSizeBytes = buffer.getAllocationSize();
          totalUsedMemory += bufferSizeBytes;
        }
      }
    }

    // Also count buffers that are not in buffersPerFifo
    for (auto &buffer : buffers) {
      // Check if this buffer is allocated on the target tile
      if (buffer.getTile() == targetTile.getResult()) {
        auto bufferSizeBytes = buffer.getAllocationSize();
        totalUsedMemory += bufferSizeBytes;
      }
    }

    return totalUsedMemory;
  }

  /// Function to analyze cross-tile buffer allocations in splitFifos
  /// Returns a simple map of (ObjectFifoCreateOp, bool) indicating cross-tile
  /// issues
  std::map<ObjectFifoCreateOp, bool>
  analyzeCrossTileFIFOBuffers(ObjectFifoState &state) {
    std::map<ObjectFifoCreateOp, bool> crossTileMap;

    for (size_t i = 0; i < state.splitFifos.size(); i++) {
      auto &[producerFifo, consumerFifos] = state.splitFifos[i];

      // Analyze producer buffers
      bool producerHasCrossTile = false;

      ObjectFifoCreateOp target = producerFifo;
      auto linkOp = getOptionalLinkOp(producerFifo);

      if (linkOp &&
          state.objFifoLinks.find(*linkOp) != state.objFifoLinks.end()) {
        target = state.objFifoLinks[*linkOp]; // Use the linked target FIFO
      }

      if (state.buffersPerFifo.find(target) != state.buffersPerFifo.end()) {
        // For each FIFO (producer and consumer):
        auto &producerBuffers = state.buffersPerFifo[target];
        TileOp expectedTile = target.getProducerTileOp();
        for (auto &buffer : producerBuffers) {
          TileOp bufferTile = buffer.getTile().getDefiningOp<TileOp>();
          if (bufferTile != expectedTile) {
            producerHasCrossTile = true;
            break;
          }
        }
      }
      crossTileMap[producerFifo] = producerHasCrossTile;

      // Analyze consumer buffers
      for (auto &consumerFifo : consumerFifos) {
        bool consumerHasCrossTile = false;
        ObjectFifoCreateOp target = consumerFifo;
        auto linkOp = getOptionalLinkOp(consumerFifo);
        if (linkOp &&
            state.objFifoLinks.find(*linkOp) != state.objFifoLinks.end()) {
          target = state.objFifoLinks[*linkOp]; // Use the linked target FIFO
        }

        if (state.buffersPerFifo.find(target) != state.buffersPerFifo.end()) {
          // For each FIFO (producer and consumer):
          auto &consumerBuffers = state.buffersPerFifo[target];
          TileOp expectedTile = target.getProducerTileOp();
          for (auto &buffer : consumerBuffers) {
            TileOp bufferTile = buffer.getTile().getDefiningOp<TileOp>();
            if (bufferTile != expectedTile) {
              consumerHasCrossTile = true;
              break;
            }
          }
        }
        crossTileMap[consumerFifo] = consumerHasCrossTile;
      }
    }
    return crossTileMap;
  }

  /// Helper function to find a tile at specific coordinates.
  /// If a tile is not found, it creates a new one and returns it.
  /// hostTile is the original tile from which we are searching for neighbors.
  /// we create the new tile below the hostTile
  TileOp findOrCreateTile(OpBuilder &builder, DeviceOp &dev, TileOp hostTile,
                          int col, int row) {
    // First, try to find an existing tile
    for (auto tile : dev.getOps<TileOp>()) {
      if (tile.getCol() == col && tile.getRow() == row) {
        return tile;
      }
    }

    // If not found, create a new one.
    OpBuilder::InsertionGuard g(builder);

    auto savedInsertionPoint = builder.saveInsertionPoint();

    // Find the last buffer operation after the host tile
    Operation *insertAfter = hostTile.getOperation();
    Operation *nextOp = insertAfter->getNextNode();
    while (nextOp && isa<BufferOp>(nextOp)) {
      insertAfter = nextOp;
      nextOp = nextOp->getNextNode();
    }

    builder.setInsertionPointAfter(insertAfter);
    auto newTile = TileOp::create(builder, builder.getUnknownLoc(), col, row);

    builder.restoreInsertionPoint(savedInsertionPoint);

    return newTile;
  }

  /// Function used to create objectFifo elements and their locks.
  /// It maps the input objectFifo to associated buffers and locks.
  void createObjectFifoElements(OpBuilder &builder, LockAnalysis &lockAnalysis,
                                ObjectFifoCreateOp op, int share_direction,
                                ObjectFifoState &state) {
    if (!op.size())
      return;

    if (op.getAieStream())
      return;

    std::vector<BufferOp> buffers;
    auto fifo = llvm::cast<AIEObjectFifoType>(op.getElemType());
    auto elemType = llvm::cast<MemRefType>(fifo.getElementType());
    int numElem = op.size();
    int of_elem_index = 0; // used to give objectFifo elements a symbolic name

    // if this objectFifo is linked to another, check if the other's elements
    // have already been created: if none of the output objectfifos of the link
    // have initValues, then the elements that are created are those of the
    // objFifo with elements of bigger size
    bool linked = false;
    auto linkOp = getOptionalLinkOp(op);
    if (linkOp) {
      auto fifoIn = linkOp->getInputObjectFifos()[0];
      auto fifoOut = linkOp->getOutputObjectFifos()[0];
      linked = true;
      if (state.objFifoLinks.find(*linkOp) != state.objFifoLinks.end())
        return; // elements have already been created
      if (linkOp->isJoin()) {
        // if join, fifoOut has bigger size
        if (op.name() != fifoOut.name())
          return;
      } else if (linkOp->isDistribute()) {
        // if distribute, fifoIn has bigger size
        if (op.name() != fifoIn.name())
          return;
      } else {
        // check if output objectfifo has initValues
        if (fifoOut.getInitValues().has_value()) {
          if (fifoOut.name() != op.name())
            return;
        } else {
          // check which objectfifo of the link has bigger size
          auto fifoInType = llvm::cast<AIEObjectFifoType>(fifoIn.getElemType());
          auto elemInType = llvm::cast<MemRefType>(fifoInType.getElementType());
          int inSize = elemInType.getNumElements();

          auto fifoOutType =
              llvm::cast<AIEObjectFifoType>(fifoOut.getElemType());
          auto elemOutType =
              llvm::cast<MemRefType>(fifoOutType.getElementType());

          if (int outSize = elemOutType.getNumElements(); inSize >= outSize) {
            if (op.name() != fifoIn.name())
              return;
          } else {
            // When output has padDimensions, MemTile buffer should use
            // input (smaller) size — padding is applied on-the-fly by DMA
            bool outHasPadding = fifoOut.getPadDimensions().has_value();
            if (outHasPadding) {
              if (op.name() != fifoIn.name())
                return;
            } else {
              if (fifoOut.name() != op.name())
                return;
            }
          }
        }
      }
    }

    TileOp creation_tile;
    auto consumerTileOp =
        dyn_cast<TileOp>(op.getConsumerTiles()[0].getDefiningOp());
    if (share_direction != 1)
      creation_tile = op.getProducerTileOp();
    else
      creation_tile = consumerTileOp;

    std::optional<ObjectFifoAllocateOp> opAlloc = getOptionalAllocateOp(op);
    if (opAlloc.has_value()) {
      TileOp delegate = opAlloc->getDelegateTileOp();
      int prodShareDir;
      int consShareDir;
      isSharedMemory(delegate, op.getProducerTileOp(), &prodShareDir);
      isSharedMemory(delegate, consumerTileOp, &consShareDir);
      if ((prodShareDir == -1 || prodShareDir == 2) &&
          (consShareDir == -1 || consShareDir == 2))
        creation_tile = delegate;
      else
        opAlloc->emitOpError("objectfifo has no shared memory access to "
                             "delegate tile's memory module");
    }

    // Reset opbuilder location to after the last tile declaration
    Operation *t = nullptr;
    auto dev = op->getParentOfType<DeviceOp>();
    for (auto tile_op : dev.getBody()->getOps<TileOp>()) {
      t = tile_op.getOperation();
    }

    builder.setInsertionPointAfter(t);
    for (int i = 0; i < numElem; i++) {

      mlir::ElementsAttr initValues = nullptr;
      if (!creation_tile.isShimTile()) {
        if (op.getInitValues().has_value()) {
          initValues =
              llvm::cast<mlir::ElementsAttr>(op.getInitValues().value()[i]);
        }

        auto elementType = elemType.getElementType();

        DataLayout dataLayout = DataLayout::closest(op.getOperation());
        int64_t elementBitWidth = dataLayout.getTypeSizeInBits(elementType);

        auto totalSizeBytes = elemType.getNumElements() * elementBitWidth / 8;
        auto &targetModel = dev.getTargetModel();

        int maxDataMemorySize = 0;
        if (creation_tile.isMemTile())
          maxDataMemorySize =
              targetModel.getMemTileSize(); // getMemTileSize returns in Bytes
        else
          maxDataMemorySize =
              targetModel
                  .getLocalMemorySize(); // getLocalMemorySize returns in Bytes

        // also need to count the buffers that are not in buffersPerFifo
        int currentUsedMemory = calculateCurrentUsedMemory(
            creation_tile, state.buffersPerFifo, buffers);

        // Check if current tile can hold the new buffer or not
        TileOp current_buf_allocation_tile =
            creation_tile; // used to keep track of the tile where the buffer is
                           // allocated
        if (creation_tile.isMemTile()) {
          if (static_cast<int>(currentUsedMemory + totalSizeBytes) >
              maxDataMemorySize) {
            // if not, check if the neighbour can hold the new buffer or not
            // Find neighbor tiles with shared memory
            std::vector<TileOp> neighborTiles;
            int currentCol = creation_tile.getCol();
            int currentRow = creation_tile.getRow();

            // Check tile to the left
            if (currentCol > 0) {
              TileOp leftTile = findOrCreateTile(builder, dev, creation_tile,
                                                 currentCol - 1, currentRow);

              int share_direction = 0;
              if (isSharedMemory(creation_tile, leftTile, &share_direction) &&
                  (share_direction == 1 || share_direction == 2)) {
                neighborTiles.push_back(leftTile);
              }
            }

            // Check tile to the right
            if (currentCol < (targetModel.columns() - 1)) {
              TileOp rightTile = findOrCreateTile(builder, dev, creation_tile,
                                                  currentCol + 1, currentRow);
              int share_direction = 0;
              if (isSharedMemory(creation_tile, rightTile, &share_direction) &&
                  (share_direction == 1 || share_direction == 2)) {
                neighborTiles.push_back(rightTile);
              }
            }

            // try to allocate on neighbor tiles
            if (!neighborTiles.empty()) {
              for (auto &tile : neighborTiles) {
                // Try to allocate on this neighbor tile
                int neighborUsedMemory = calculateCurrentUsedMemory(
                    tile, state.buffersPerFifo, buffers);
                if (static_cast<int>(neighborUsedMemory + totalSizeBytes) <=
                    maxDataMemorySize) {
                  // Allocate buffer on neighbor tile, change creation_tile to
                  // be this neighbour tile
                  current_buf_allocation_tile = tile;
                  break;
                }
              }
            }
          }
        }
        auto buff = BufferOp::create(
            builder, builder.getUnknownLoc(), elemType,
            current_buf_allocation_tile,
            builder.getStringAttr(op.name().str() + "_buff_" +
                                  std::to_string(of_elem_index)),
            /*address*/ nullptr, initValues,
            /*mem_bank*/ nullptr);
        buffers.push_back(buff);
      }
      of_elem_index++;
    }

    int repeatCount = 1;
    int joinDistribFactor = 1;
    if (op.getRepeatCount().has_value())
      repeatCount = op.getRepeatCount().value();
    if (linked) {
      if (linkOp->getRepeatCount().has_value())
        repeatCount = linkOp->getRepeatCount().value();
      if (linkOp->isDistribute())
        joinDistribFactor *= linkOp->getFifoOuts().size();
      else if (linkOp->isJoin())
        joinDistribFactor *= linkOp->getFifoIns().size();
      state.objFifoLinks[*linkOp] = op;
    }
    std::vector<LockOp> locks = createObjectFifoLocks(
        builder, lockAnalysis, op, numElem, joinDistribFactor, creation_tile,
        repeatCount, state);
    state.buffersPerFifo[op] = buffers;
    state.locksPerFifo[op] = locks;
  }

  /// Function that returns a pointer to the block of a Region
  /// that contains the AIEEndOp.
  Block *findEndOpBlock(Region &r) {
    Block *endBlock = nullptr;
    for (auto &bl : r.getBlocks())
      if (!bl.getOps<EndOp>().empty())
        endBlock = &bl;
    return endBlock;
  }

  /// Function used to create a Bd block.
  ///
  /// Returns the newly created DMABDOp so the caller can decorate it with
  /// dataflow-source-specific discardable attributes (e.g. the
  /// SparseFifo (de)compression bit propagated from the originating
  template <typename MyOp>
  DMABDOp createBd(OpBuilder &builder, LockOp acqLock, int acqMode,
                   LockAction acqLockAction, LockOp relLock, int relMode,
                   MyOp buff, int offset, int len, Block *succ,
                   BDDimLayoutArrayAttr dims,
                   BDPadLayoutArrayAttr padDimensions,
                   std::optional<PacketInfoAttr> bdPacket) {
    if (acqLock)
      UseLockOp::create(builder, builder.getUnknownLoc(), acqLock,
                        acqLockAction, acqMode);
    if (bdPacket) {
      DMABDPACKETOp::create(builder, builder.getUnknownLoc(),
                            bdPacket->getPktType(), bdPacket->getPktId());
    }
    DMABDOp bdOp = [&]() {
      if (!dims.getValue().empty() && padDimensions) {
        return DMABDOp::create(builder, builder.getUnknownLoc(), buff, offset,
                               len, dims, padDimensions);
      }
      if (!dims.getValue().empty()) {
        return DMABDOp::create(builder, builder.getUnknownLoc(), buff, offset,
                               len, dims);
      }
      return DMABDOp::create(builder, builder.getUnknownLoc(), buff, offset,
                             len);
    }();
    if (acqLock)
      UseLockOp::create(builder, builder.getUnknownLoc(), relLock,
                        LockAction::Release, relMode);
    NextBDOp::create(builder, builder.getUnknownLoc(), succ);
    return bdOp;
  }

  /// Function used to create a Bd block.
  /// If lockMode is 0 we create a consumerDMA (i.e. on producer tile) else a
  /// producerDMA (i.e. on consumer tile).
  template <typename MyOp>
  void createBdBlock(OpBuilder &builder, ObjectFifoCreateOp op, int lockMode,
                     int acqNum, int relNum, MyOp buff, int offset, int len,
                     DMAChannelDir channelDir, size_t lockIndex, Block *succ,
                     BDDimLayoutArrayAttr dims,
                     BDPadLayoutArrayAttr padDimensions,
                     std::optional<PacketInfoAttr> bdPacket,
                     ObjectFifoState &state, bool distribOrJoin = false) {
    LockOp acqLock;
    LockOp relLock;
    int acqMode = 1;
    int relMode = 1;
    auto acqLockAction = LockAction::Acquire;
    if (state.locksPerFifo[op].size() > 0) {
      auto dev = op->getParentOfType<DeviceOp>();
      if (auto &target = dev.getTargetModel();
          target.getTargetArch() == AIEArch::AIE1) {
        acqMode = lockMode == 0 ? 1 : 0;
        relMode = lockMode == 0 ? 0 : 1;
        acqLock = state.locksPerFifo[op][lockIndex];
        relLock = state.locksPerFifo[op][lockIndex];
      } else {
        acqMode = acqNum;
        relMode = relNum;
        acqLockAction = LockAction::AcquireGreaterEqual;
        int prodLockIndex = 0;
        int consLockIndex = 1;
        if (distribOrJoin) {
          prodLockIndex = lockIndex * 2;
          consLockIndex = lockIndex * 2 + 1;
        }
        acqLock = channelDir == DMAChannelDir::S2MM
                      ? state.locksPerFifo[op][prodLockIndex]
                      : state.locksPerFifo[op][consLockIndex];
        relLock = channelDir == DMAChannelDir::S2MM
                      ? state.locksPerFifo[op][consLockIndex]
                      : state.locksPerFifo[op][prodLockIndex];
      }
    }
    DMABDOp bdOp = createBd(builder, acqLock, acqMode, acqLockAction, relLock,
                            relMode, buff, offset, len, succ, dims,
                            padDimensions, bdPacket);

    // originating ObjectFifoCreateOp onto each DMABDOp we just created so
    // downstream BD-emit (AIEDMATasksToNPU -> AIEDmaToNpu) can flip the
    // per-channel ``Enable_Compression`` bit on the AIE2/AIE2P tile DMA BD
    // config word. See ``python/iron/sparse.py`` for the discardable-attr
    // contract; see AM020 Ch. 2 p. 27 for the hardware bit. The default
    // (no attrs on the ObjectFifoCreateOp) is the pre-existing behaviour:
    // no discardable attr on the DMABDOp -> no compression bit set.
    propagateSparseCompressionAttr(bdOp.getOperation(), op.getOperation(),
                                   channelDir);
  }

  /// discardable attrs (set by ``aie.iron.sparse.SparseFifo.resolve``)
  /// from the originating ObjectFifoCreateOp and, if the channel
  /// direction selects the matching half of the
  /// (compress_mm2s, decompress_s2mm) pair AND the BD lives on a
  /// compute (AIE) tile, attach a discardable boolean
  /// ``aie.enable_compression = true`` attribute on the new DMABDOp.
  /// the cross-module footgun guard): the SparseFifo lowering
  /// already attaches the intent to the ObjectFifoCreateOp; this
  /// propagates it to the DMABDOp where the ObjectFifoCreateOp
  /// itself is about to be erased.
  ///
  /// Cross-module footgun (AM029 / aie_registers_aie2.json):
  ///   * Compute-tile MEMORY_MODULE DMA_BD0_1 bit 31 = Enable_Compression
  ///   * Memory-tile MEMORY_TILE_MODULE DMA_BD0_1 bits 31:26 = D0_Pad_Before
  ///   * Shim DMA: AM029 documents Enable_Compression for "AIE-ML
  ///     memory and AIE-ML tile DMA" only — not for shim.
  /// Setting aie.enable_compression on a memtile or shim BD would
  /// either silently corrupt an unrelated field (memtile) or set an
  /// undocumented bit (shim). So we walk up to the BD's owning tile
  /// and bail out unless it's a compute tile.
  static void propagateSparseCompressionAttr(Operation *bdOp, Operation *fifoOp,
                                             DMAChannelDir channelDir) {
    if (!bdOp || !fifoOp)
      return;
    StringRef attrName;
    if (channelDir == DMAChannelDir::MM2S) {
      attrName = "aie.compress_mm2s";
    } else if (channelDir == DMAChannelDir::S2MM) {
      attrName = "aie.decompress_s2mm";
    } else {
      return;
    }
    auto enable = fifoOp->getAttrOfType<BoolAttr>(attrName);
    if (!enable || !enable.getValue())
      return;

    // Cross-module footgun guard: only emit on compute-tile BDs. The
    // BD lives in either MemOp (compute), MemTileDMAOp (memtile), or
    // ShimDMAOp (shim). Walk up to find the parent and check the
    // tile type via the device's target model.
    TileOp tileOp;
    if (auto memOp = bdOp->getParentOfType<MemOp>())
      tileOp = memOp.getTileOp();
    else if (bdOp->getParentOfType<MemTileDMAOp>() ||
             bdOp->getParentOfType<ShimDMAOp>())
      return; // memtile / shim — bit means something else (or undocumented).
    else
      return; // unknown parent; conservative skip.

    if (!tileOp)
      return;
    auto deviceOp = tileOp->getParentOfType<DeviceOp>();
    if (!deviceOp)
      return;
    const auto &targetModel = deviceOp.getTargetModel();
    if (tileOp.isShimTile() ||
        targetModel.isMemTile(tileOp.getCol(), tileOp.getRow()))
      return;

    bdOp->setAttr("aie.enable_compression",
                  BoolAttr::get(bdOp->getContext(), true));
  }

  /// Function that either calls createAIETileDMA(), createShimDMA() or
  /// createMemTileDMA() based on op tile row value.
  void createDMA(DeviceOp &device, OpBuilder &builder, ObjectFifoCreateOp op,
                 DMAChannelDir channelDir, int channelIndex, int lockMode,
                 BDDimLayoutArrayAttr dims, BDPadLayoutArrayAttr pad_dims,
                 std::optional<PacketInfoAttr> bdPacket,
                 ObjectFifoState &state) {
    if (op.getProducerTileOp().isShimTile()) {
      createShimDMA(device, builder, op, channelDir, channelIndex, lockMode,
                    dims, bdPacket, state);
    } else if (op.getProducerTileOp().isMemTile()) {
      BDPadLayoutArrayAttr padDims = nullptr;
      if (channelDir == DMAChannelDir::MM2S && pad_dims)
        padDims = pad_dims;
      createMemTileDMA(device, builder, op, channelDir, channelIndex, lockMode,
                       dims, padDims, bdPacket, state);
    } else {
      createAIETileDMA(device, builder, op, channelDir, channelIndex, lockMode,
                       dims, bdPacket, state);
    }
  }

  /// Function used to create a MemOp region with a DMA channel.
  /// It uses creatBdBlock(), see there for lockMode input.
  void createAIETileDMA(DeviceOp &device, OpBuilder &builder,
                        ObjectFifoCreateOp op, DMAChannelDir channelDir,
                        int channelIndex, int lockMode,
                        BDDimLayoutArrayAttr dims,
                        std::optional<PacketInfoAttr> bdPacket,
                        ObjectFifoState &state) {
    size_t numBlocks = op.size();
    if (numBlocks == 0)
      return;

    int acqNum = 1;
    int relNum = 1;

    auto fifo = llvm::cast<AIEObjectFifoType>(op.getElemType());
    auto elemType = llvm::cast<MemRefType>(fifo.getElementType());
    int len = elemType.getNumElements();

    // check for repeat count
    int repeatCount = 1;
    if (op.getRepeatCount().has_value())
      repeatCount = op.getRepeatCount().value();

    // search for the buffers/locks (based on if this objFifo has a link)
    ObjectFifoCreateOp target = op;
    if (std::optional<ObjectFifoLinkOp> linkOp = getOptionalLinkOp(op);
        linkOp.has_value()) {
      if (state.objFifoLinks.find(linkOp.value()) != state.objFifoLinks.end()) {
        target = state.objFifoLinks[linkOp.value()];
        if (target == op) {
          if (linkOp->getRepeatCount().has_value()) {
            acqNum *= linkOp->getRepeatCount().value();
            relNum *= linkOp->getRepeatCount().value();
          }
        }
      }
    }

    // search for MemOp
    Operation *producerMem = nullptr;
    for (auto memOp : device.getOps<MemOp>()) {
      if (memOp.getTile() == op.getProducerTile()) {
        producerMem = memOp.getOperation();
        break;
      }
    }

    // if none exists, create one
    TileOp objFifoTileOp = target.getProducerTileOp();
    if (producerMem == nullptr) {
      OpBuilder::InsertionGuard g(builder);
      builder.setInsertionPoint(device.getBody()->getTerminator());
      auto newMemOp =
          MemOp::create(builder, builder.getUnknownLoc(), objFifoTileOp);
      {
        OpBuilder::InsertionGuard g(builder);
        builder.setInsertionPointToStart(&newMemOp.getRegion().emplaceBlock());
        EndOp::create(builder, builder.getUnknownLoc());
      }
      producerMem = newMemOp.getOperation();
    }
    Block *endBlock = findEndOpBlock(producerMem->getRegion(0));
    Block *lastDmaBlock = endBlock->getSinglePredecessor();
    Block *dmaBlock = builder.createBlock(endBlock);
    Block *bdBlock = builder.createBlock(endBlock);

    // create DMA channel
    builder.setInsertionPointToStart(dmaBlock);
    DMAStartOp::create(builder, builder.getUnknownLoc(), channelDir,
                       channelIndex, /*repeatCout*/ 0, bdBlock, endBlock);
    if (lastDmaBlock != nullptr)
      lastDmaBlock->getTerminator()->setSuccessor(dmaBlock, 1);

    // create Bd blocks
    Block *succ;
    Block *curr = bdBlock;
    size_t elemIndex = 0;
    size_t totalBlocks = 0;
    for (size_t i = 0; i < numBlocks; i++) {
      if (elemIndex >= state.buffersPerFifo[target].size())
        break;
      for (int r = 0; r < repeatCount; r++) {
        if (totalBlocks == numBlocks * repeatCount - 1)
          succ = bdBlock;
        else
          succ = builder.createBlock(endBlock);

        builder.setInsertionPointToStart(curr);
        createBdBlock<BufferOp>(builder, target, lockMode, acqNum, relNum,
                                state.buffersPerFifo[target][elemIndex],
                                /*offset*/ 0, len, channelDir, elemIndex, succ,
                                dims, nullptr, bdPacket, state);
        curr = succ;
        totalBlocks++;
      }
      elemIndex++;
    }
  }

  /// Function used to create a ShimDMAOp region with a DMA channel.
  /// It uses creatBdBlock(), see there for lockMode input.
  void createShimDMA(DeviceOp &device, OpBuilder &builder,
                     ObjectFifoCreateOp op, DMAChannelDir channelDir,
                     int channelIndex, int lockMode, BDDimLayoutArrayAttr dims,
                     std::optional<PacketInfoAttr> bdPacket,
                     ObjectFifoState &state) {
    size_t numBlocks = state.externalBuffersPerFifo[op].size();
    if (numBlocks == 0)
      return;

    int acqNum = 1;
    int relNum = 1;

    // search for ShimDMAOp
    Operation *producerDMA = nullptr;
    for (auto dmaOp : device.getOps<ShimDMAOp>()) {
      if (dmaOp.getTile() == op.getProducerTile()) {
        producerDMA = dmaOp.getOperation();
        break;
      }
    }

    // if none exists, create one
    TileOp objFifoTileOp = op.getProducerTileOp();
    if (producerDMA == nullptr) {
      OpBuilder::InsertionGuard g(builder);
      builder.setInsertionPoint(device.getBody()->getTerminator());
      auto newDMAOp = ShimDMAOp::create(builder, builder.getUnknownLoc(),
                                        builder.getIndexType(), objFifoTileOp);
      {
        OpBuilder::InsertionGuard g(builder);
        builder.setInsertionPointToStart(&newDMAOp.getRegion().emplaceBlock());
        EndOp::create(builder, builder.getUnknownLoc());
      }
      producerDMA = newDMAOp.getOperation();
    }

    Block *endBlock = findEndOpBlock(producerDMA->getRegion(0));
    Block *lastDmaBlock = endBlock->getSinglePredecessor();
    Block *dmaBlock = builder.createBlock(endBlock);
    Block *bdBlock = builder.createBlock(endBlock);

    // create DMA channel
    builder.setInsertionPointToStart(dmaBlock);
    DMAStartOp::create(builder, builder.getUnknownLoc(), channelDir,
                       channelIndex, /*repeatCout*/ 0, bdBlock, endBlock);
    if (lastDmaBlock != nullptr)
      lastDmaBlock->getTerminator()->setSuccessor(dmaBlock, 1);

    // create Bd blocks
    Block *succ;
    Block *curr = bdBlock;
    size_t elemIndex = 0;
    for (size_t i = 0; i < numBlocks; i++) {
      if (elemIndex >= state.externalBuffersPerFifo[op].size())
        break;
      if (i == numBlocks - 1)
        succ = bdBlock;
      else
        succ = builder.createBlock(endBlock);

      MemRefType buffer = state.externalBuffersPerFifo[op][elemIndex].getType();
      int len = buffer.getNumElements();
      builder.setInsertionPointToStart(curr);
      createBdBlock<ExternalBufferOp>(
          builder, op, lockMode, acqNum, relNum,
          state.externalBuffersPerFifo[op][elemIndex],
          /*offset*/ 0, len, channelDir, elemIndex, succ, dims, nullptr,
          bdPacket, state);
      curr = succ;
      elemIndex++;
    }
  }

  /// Function used to create a MemTileDMAOp region with a DMA channel.
  /// It uses creatBdBlock(), see there for lockMode input.
  void createMemTileDMA(DeviceOp &device, OpBuilder &builder,
                        ObjectFifoCreateOp op, DMAChannelDir channelDir,
                        int channelIndex, int lockMode,
                        BDDimLayoutArrayAttr dims,
                        BDPadLayoutArrayAttr padDimensions,
                        std::optional<PacketInfoAttr> bdPacket,
                        ObjectFifoState &state) {
    size_t numBlocks = op.size();
    if (numBlocks == 0)
      return;

    auto fifo = llvm::cast<AIEObjectFifoType>(op.getElemType());
    auto elemType = llvm::cast<MemRefType>(fifo.getElementType());
    int lenOut = elemType.getNumElements();
    int acqNum = 1;
    int relNum = 1;

    // check for repeat count
    int repeatCount = 1;
    if (op.getRepeatCount().has_value())
      repeatCount = op.getRepeatCount().value();

    // check for BD chain repeat count
    auto bdChainIterCount = op.getIterCount();

    // search for the buffers/locks (based on if this objFifo has a link)
    // identify size difference between input and output memrefs
    ObjectFifoCreateOp target = op;
    bool isDistribute = false;
    bool isJoin = false;
    int extraOffset = 0;
    int joinDistribFactor = 1;
    int joinDistribLockIndex = 0;
    auto linkOp = getOptionalLinkOp(op);
    if (linkOp) {
      if (state.objFifoLinks.find(*linkOp) != state.objFifoLinks.end()) {
        target = state.objFifoLinks[*linkOp];
        auto srcOffsets = linkOp->getSrcOffsets();
        auto dstOffsets = linkOp->getDstOffsets();

        if (linkOp->getRepeatCount().has_value())
          if (linkOp->getInputObjectFifos()[0] == op) {
            acqNum *= linkOp->getRepeatCount().value();
            relNum *= linkOp->getRepeatCount().value();
          }

        if (linkOp->isJoin()) {
          // compute offset and length
          isJoin = true;
          if (target == op) {
            joinDistribFactor *= linkOp->getFifoIns().size();
          } else {
            int i = 0;
            for (auto fifoIn : linkOp->getInputObjectFifos()) {
              if (fifoIn.name() == op.name())
                break;
              i++;
            }
            extraOffset = *getConstantIntValue(srcOffsets[i]);
            lenOut = linkOp->getJoinTransferLengths()[i];
            joinDistribLockIndex = i;
          }
        } else if (linkOp->isDistribute()) {
          // compute offset and length
          isDistribute = true;
          if (target == op) {
            joinDistribFactor *= linkOp->getFifoOuts().size();
          } else {
            int i = 0;
            for (auto fifoOut : linkOp->getOutputObjectFifos()) {
              if (fifoOut.name() == op.name())
                break;
              i++;
            }
            extraOffset = *getConstantIntValue(dstOffsets[i]);
            lenOut = linkOp->getDistributeTransferLengths()[i];
            joinDistribLockIndex = i;
          }
        } else {
          if (target != op) {
            auto targetFifo =
                llvm::cast<AIEObjectFifoType>(target.getElemType());
            auto targetElemType =
                llvm::cast<MemRefType>(targetFifo.getElementType());
            int targetLen = targetElemType.getNumElements();
            // Only override when target is larger or equal. When target
            // is smaller (padDimensions size mismatch after buffer
            // ownership change), op's own element count is correct.
            if (targetLen >= lenOut)
              lenOut = targetLen;
          }
        }

        // check if current op is of smaller size in link
        if (target != op) {
          numBlocks = target.size();
        }
      }
    }

    // search for MemTileDMAOp
    Operation *producerDMA = nullptr;
    for (auto dmaOp : device.getOps<MemTileDMAOp>()) {
      if (dmaOp.getTile() == target.getProducerTile()) {
        producerDMA = dmaOp.getOperation();
        break;
      }
    }

    // if none exists, create one
    TileOp objFifoTileOp = target.getProducerTileOp();
    if (producerDMA == nullptr) {
      OpBuilder::InsertionGuard g(builder);
      builder.setInsertionPoint(device.getBody()->getTerminator());
      auto newDMAOp =
          MemTileDMAOp::create(builder, builder.getUnknownLoc(), objFifoTileOp);
      {
        OpBuilder::InsertionGuard g(builder);
        builder.setInsertionPointToStart(&newDMAOp.getRegion().emplaceBlock());
        EndOp::create(builder, builder.getUnknownLoc());
      }
      producerDMA = newDMAOp.getOperation();
    }

    Block *endBlock = findEndOpBlock(producerDMA->getRegion(0));
    Block *lastDmaBlock = endBlock->getSinglePredecessor();
    Block *dmaBlock = builder.createBlock(endBlock);
    Block *bdBlock = builder.createBlock(endBlock);

    // create DMA channel
    builder.setInsertionPointToStart(dmaBlock);

    // Use iter_count if available, otherwise default to 0
    int taskCount = 0;
    bool isBdChainMode = false;
    if (bdChainIterCount.has_value()) {
      taskCount = bdChainIterCount.value() - 1;
      isBdChainMode = true;
    }
    DMAStartOp::create(builder, builder.getUnknownLoc(), channelDir,
                       channelIndex, taskCount, bdBlock, endBlock);
    if (lastDmaBlock != nullptr)
      lastDmaBlock->getTerminator()->setSuccessor(dmaBlock, 1);

    // create Bd blocks
    Block *succ;
    Block *curr = bdBlock;
    size_t elemIndex = 0;
    size_t lockIndex = 0;
    size_t totalBlocks = 0;
    bool distribOrJoin = false;

    for (size_t i = 0; i < numBlocks; i++) {
      if (elemIndex >= state.buffersPerFifo[target].size())
        break;
      for (int r = 0; r < repeatCount * joinDistribFactor; r++) {
        if (totalBlocks == numBlocks * repeatCount * joinDistribFactor - 1) {
          // If iter_count attribute is set (BD chain mode), create a
          // dedicated terminating block
          if (isBdChainMode) {
            succ = builder.createBlock(endBlock);
            // Create a separate terminating block with aie.end for this
            // specific DMA channel
            builder.setInsertionPointToStart(succ);
            EndOp::create(builder, builder.getUnknownLoc());
          } else {
            succ = bdBlock;
          }
        } else {
          succ = builder.createBlock(endBlock);
        }

        builder.setInsertionPointToStart(curr);
        int offset = 0;
        if (isDistribute || isJoin) {
          distribOrJoin = true;
          if (target == op) {
            if (isDistribute) {
              offset = *getConstantIntValue(linkOp->getDstOffsets()[r]);
              lenOut = linkOp->getDistributeTransferLengths()[r];
            } else {
              offset = *getConstantIntValue(linkOp->getSrcOffsets()[r]);
              lenOut = linkOp->getJoinTransferLengths()[r];
            }
            lockIndex = r % joinDistribFactor;
          } else {
            offset = extraOffset;
            lockIndex = joinDistribLockIndex;
          }
        } else {
          lockIndex = elemIndex;
        }

        createBdBlock<BufferOp>(builder, target, lockMode, acqNum, relNum,
                                state.buffersPerFifo[target][elemIndex], offset,
                                lenOut, channelDir, lockIndex, succ, dims,
                                padDimensions, bdPacket, state, distribOrJoin);
        curr = succ;
        totalBlocks++;
      }
      elemIndex++;
    }
  }

  // Function that computes the Least Common Multiplier of the values
  // of a vector.
  int computeLCM(std::set<int> values) {
    int lcm = 1;
    for (int i : values)
      lcm = i * lcm / std::gcd(i, lcm);
    return lcm;
  }

  // Function that unrolls for-loops that contain objectFifo operations.
  LogicalResult unrollForLoops(DeviceOp &device, OpBuilder &builder,
                               std::set<TileOp> objectFifoTiles) {
    for (auto coreOp : device.getOps<CoreOp>()) {
      if (objectFifoTiles.count(coreOp.getTileOp()) > 0) {
        std::vector<scf::ForOp> unrolledLoops;
        std::map<Operation *, bool> foundMap;
        std::map<Operation *, int64_t> remainderMap;
        std::map<Operation *, int64_t> tripCountMap;
        WalkResult res = coreOp.walk([&](scf::ForOp forLoop) {
          // look for operations on objectFifos
          // when multiple fifos in same loop, must use the smallest
          // common multiplier as the unroll factor
          foundMap[forLoop.getOperation()] = false;
          std::set<int> objFifoSizes;
          Block *body = forLoop.getBody();
          remainderMap[forLoop.getOperation()] = 0;
          for (auto acqOp : body->getOps<ObjectFifoAcquireOp>()) {
            if (acqOp.getOperation()->getParentOp() == forLoop) {
              ObjectFifoCreateOp op = acqOp.getObjectFifo();
              // VariableRateFifo opts out of LCM-based loop unrolling.
              // The producer's loop body contains a conditional
              // acquire/release that the LCM-unroll math cannot model;
              // the runtime-counter machinery handles asymmetric rates
              // correctly without unrolling. If the loop has only
              // variable-rate accesses, foundMap stays false and the
              // loop is left alone.
              auto vrAttr = op->getAttrOfType<BoolAttr>("aie.variable_rate");
              if (vrAttr && vrAttr.getValue()) {
                continue;
              }
              foundMap[forLoop.getOperation()] = true;
              objFifoSizes.insert(op.size());
            }
          }
          // If the loop doesn't have acquire and release locks
          // Push it to the unrolledLoops to avoid unrolling
          if (!foundMap[forLoop.getOperation()]) {
            unrolledLoops.push_back(forLoop);
            return WalkResult::advance();
          }
          // Walk in the loop region to unroll the loop and its remainder
          Region *region = forLoop->getParentRegion();
          scf::ForOp prevLoop;
          prevLoop = forLoop;
          tripCountMap[prevLoop.getOperation()] = 0;
          while (remainderMap[prevLoop.getOperation()] > 1 ||
                 foundMap[prevLoop.getOperation()]) {
            region->walk([&](scf::ForOp remLoop) {
              bool skipLoop = false;
              int64_t tripCount = 0;
              if (remLoop.getSingleLowerBound() &&
                  remLoop.getSingleUpperBound() && remLoop.getSingleStep()) {
                tripCount = remLoop.getStaticTripCount()->getSExtValue();
              }
              int unrollFactor =
                  computeLCM(objFifoSizes); // also counts original loop body
              // Loop ids are not unique.
              // Sometimes, immediately after unrolling, the unrolled loop
              // and the one next to it (can be the remainder loop or an
              // independent loop) will have the same ID. This makes it
              // difficult to identify which loop needs to be unrolled.
              // Once it restarts walking from start, it ends up allocating
              // new ID to each loop.
              if (remainderMap[prevLoop.getOperation()] > 1 &&
                  foundMap[remLoop.getOperation()] == false &&
                  prevLoop != remLoop) {
                skipLoop = true;
              }
              if (std::count(unrolledLoops.begin(), unrolledLoops.end(),
                             remLoop) == 0 &&
                  !skipLoop) {
                tripCountMap[remLoop.getOperation()] = tripCount;
                // if loop iterations < unrollFactor, unroll the loop fully
                if (tripCountMap[remLoop.getOperation()] < unrollFactor)
                  unrollFactor = tripCountMap[remLoop.getOperation()];
                // If unrollFactor = 0,divide by zero
                if (unrollFactor == 0) {
                  remLoop.emitOpError()
                      << "could not be unrolled with unrollFactor = 0, check "
                         "loop boundaries."
                      << "\n";
                  return WalkResult::interrupt();
                }
                remainderMap[remLoop.getOperation()] =
                    tripCountMap[remLoop.getOperation()] % unrollFactor;
                auto step = remLoop.getStep()
                                .getDefiningOp<arith::ConstantOp>()
                                .getValue();
                int64_t step_value = llvm::dyn_cast<IntegerAttr>(step).getInt();

                if (step_value < unrollFactor ||
                    foundMap[remLoop.getOperation()]) {
                  // Process the for loop
                  if (failed(mlir::loopUnrollByFactor(remLoop, unrollFactor))) {
                    remLoop.emitOpError()
                        << "could not be unrolled with unrollFactor: "
                        << unrollFactor << "\n";
                    return WalkResult::interrupt();
                  }
                  unrolledLoops.push_back(remLoop);
                  foundMap[remLoop.getOperation()] = false;
                } else {
                  remainderMap[remLoop.getOperation()] = 0;
                  foundMap[remLoop.getOperation()] = false;
                }
              } else {
                remainderMap[remLoop.getOperation()] = 0;
                foundMap[remLoop.getOperation()] = false;
              }
              prevLoop = remLoop;
              return WalkResult::advance();
            });
          }
          return WalkResult::advance();
        });
        if (res.wasInterrupted())
          return failure();
      }
    }
    return success();
  }

  // Function that generates the IR to update runtime state of objectfifo
  // accesses. Called by dynamicGlobalObjectFifos().
  void updateGlobalNextIndex(OpBuilder &builder, ObjectFifoReleaseOp relOp,
                             BufferOp globalNextIndex, arith::ConstantOp index,
                             arith::ConstantOp size) {
    builder.setInsertionPointAfter(relOp);
    Value oldCounter = memref::LoadOp::create(
        builder, builder.getUnknownLoc(), globalNextIndex,
        ValueRange(ArrayRef({index.getResult()})));
    Value val =
        arith::ConstantOp::create(builder, oldCounter.getLoc(),
                                  builder.getI32IntegerAttr(relOp.getSize()));
    Value sum = arith::AddIOp::create(builder, val.getLoc(), oldCounter, val);
    Value isGreaterEqual = arith::CmpIOp::create(
        builder, sum.getLoc(), arith::CmpIPredicate::sge, sum, size);
    Value newCounter = arith::SelectOp::create(
        builder, sum.getLoc(), isGreaterEqual,
        arith::SubIOp::create(builder, sum.getLoc(), sum, size), sum);
    memref::StoreOp::create(builder, size.getLoc(), newCounter, globalNextIndex,
                            ValueRange(ArrayRef({index.getResult()})));
  }

  // Function that generates the IR for objectfifo accesses to be handled at
  // runtime.
  LogicalResult dynamicGlobalObjectFifos(DeviceOp &device, OpBuilder &builder,
                                         std::set<TileOp> objectFifoTiles,
                                         ObjectFifoState &state) {
    for (auto coreOp : device.getOps<CoreOp>()) {
      if (objectFifoTiles.count(coreOp.getTileOp()) <= 0)
        continue;
      if (objectFifoTiles.count(coreOp.getTileOp()) > 0) {
        // For each core: count the number of objectFifos and create
        // a global buffer just before the core to track index of
        // next object to access.
        // !! NOTE !! objectFifos with same producer / consumer tile
        // need two counters (accessed based on the ObjectFifoPort)
        std::map<std::pair<ObjectFifoCreateOp, ObjectFifoPort>, int> fifoSizes;
        // Also, keep a map of the ConstantOps for the indices per OF
        // and a map with the ConstantOps for the sizes per OF.
        std::map<std::pair<ObjectFifoCreateOp, ObjectFifoPort>,
                 arith::ConstantOp>
            globalIndices;
        std::map<std::pair<ObjectFifoCreateOp, ObjectFifoPort>,
                 arith::ConstantOp>
            constantSizes;

        int index = 0;
        builder.setInsertionPointToStart(&(coreOp.getBody().front()));
        Value initVal = arith::ConstantOp::create(
            builder, builder.getUnknownLoc(), builder.getI32IntegerAttr(0));
        coreOp.walk([&](ObjectFifoAcquireOp acqOp) {
          ObjectFifoCreateOp op = acqOp.getObjectFifo();
          ObjectFifoPort port = acqOp.getPort();
          if (fifoSizes.find({op, port}) == fifoSizes.end()) {
            fifoSizes[{op, port}] = op.size();
            auto indexOp = arith::ConstantOp::create(
                builder, initVal.getLoc(), builder.getIndexAttr(index));
            globalIndices[{op, port}] = indexOp;
            index++;
            auto size =
                arith::ConstantOp::create(builder, indexOp.getLoc(),
                                          builder.getI32IntegerAttr(op.size()));
            constantSizes[{op, port}] = size;
          }
        });
        builder.setInsertionPoint(coreOp);
        auto memrefTy =
            MemRefType::get(SmallVector<int64_t>{(int64_t)fifoSizes.size()},
                            builder.getI32Type());
        auto globalNextIndex = BufferOp::create(
            builder, builder.getUnknownLoc(), memrefTy, coreOp.getTile(),
            /*sym_name*/ nullptr, /*address*/ nullptr,
            /*initial_value*/ nullptr, /*mem_bank*/ nullptr);

        // Initialize all counters in the global buffers to 0.
        for (auto i : constantSizes) {
          builder.setInsertionPointAfter(i.second);
          memref::StoreOp::create(
              builder, builder.getUnknownLoc(), initVal, globalNextIndex,
              ValueRange(ArrayRef({globalIndices[i.first].getResult()})));
        }

        // Walk the code:
        // - after each ObjectFifoReleaseOp:
        //    - globalNextIndex: add #rel modulo objfifo depth
        // - before each ObjectFifoAcquireOp:
        //    - globalNextIndex: load index and use it to index_switch (one
        //    IndexSwithOp per AccessOp)
        WalkResult res = coreOp.walk([&](Operation *op) {
          if (auto relOp = dyn_cast<ObjectFifoReleaseOp>(op)) {
            ObjectFifoCreateOp createOp = relOp.getObjectFifo();
            ObjectFifoPort port = relOp.getPort();
            updateGlobalNextIndex(builder, relOp, globalNextIndex,
                                  globalIndices[{createOp, port}],
                                  constantSizes[{createOp, port}]);
          }
          if (auto acqOp = dyn_cast<ObjectFifoAcquireOp>(op)) {
            std::vector<ObjectFifoSubviewAccessOp> accessOps;
            for (auto u : acqOp->getUsers())
              if (auto accessOp = dyn_cast<ObjectFifoSubviewAccessOp>(u))
                accessOps.push_back(accessOp);

            for (auto accessOp : accessOps) {
              ObjectFifoCreateOp createOp = acqOp.getObjectFifo();
              ObjectFifoPort port = acqOp.getPort();

              // Single switch case
              if (fifoSizes[{createOp, port}] == 1)
                return WalkResult::advance();

              // Create a switch for each subview access
              builder.setInsertionPointAfter(accessOp);
              auto switchIndexAsInteger = memref::LoadOp::create(
                  builder, builder.getUnknownLoc(), globalNextIndex,
                  ValueRange(
                      ArrayRef({globalIndices[{createOp, port}].getResult()})));
              auto switchIndex = arith::IndexCastOp::create(
                  builder, builder.getUnknownLoc(), builder.getIndexType(),
                  switchIndexAsInteger);
              unsigned caseRegionCounts = fifoSizes[{createOp, port}];
              SmallVector<int64_t, 4> caseValues;
              for (int i = 0; i < fifoSizes[{createOp, port}]; ++i) {
                caseValues.push_back(i);
              }
              auto cases =
                  DenseI64ArrayAttr::get(builder.getContext(), caseValues);
              auto switchOp = scf::IndexSwitchOp::create(
                  builder, switchIndex.getLoc(),
                  TypeRange({state.buffersPerFifo[createOp][0].getType()}),
                  switchIndex, cases, caseRegionCounts);
              // Create default case of IndexSwitchOp
              builder.createBlock(&switchOp.getDefaultRegion());
              auto bufferIndex = (accessOp.getIndex()) % createOp.size();
              builder.setInsertionPointToStart(&(switchOp.getDefaultBlock()));
              scf::YieldOp::create(
                  builder, builder.getUnknownLoc(),
                  state.buffersPerFifo[createOp][bufferIndex].getResult());
              for (int i = 0; i < fifoSizes[{createOp, port}]; ++i) {
                // Create other cases of IndexSwitchOp
                builder.createBlock(&switchOp.getCaseRegions()[i]);
                builder.setInsertionPoint(&switchOp.getCaseBlock(i),
                                          switchOp.getCaseBlock(i).begin());
                int bufferToBeAccesed =
                    (accessOp.getIndex() + i) % fifoSizes[{createOp, port}];
                scf::YieldOp::create(
                    builder, switchOp.getCaseRegions()[i].getLoc(),
                    state.buffersPerFifo[createOp][bufferToBeAccesed]
                        .getResult());
              }

              // Replace all uses of accessed objectfifo buffers with
              // results of switchOps
              accessOp.getOutput().replaceAllUsesWith(switchOp.getResult(0));
            }
          }
          return WalkResult::advance();
        });
        if (res.wasInterrupted())
          return failure();
      }
    }
    return success();
  }

  /// Function used to create a UseLockOp based on input parameters.
  /// acc is an accumulator map that tracks the indices of the next locks to
  /// acquire (or release). Uses op to find index of acc for next lockID.
  /// Updates acc.
  void createUseLocks(OpBuilder &builder, ObjectFifoCreateOp op,
                      ObjectFifoPort port,
                      DenseMap<std::pair<ObjectFifoCreateOp, int>, int> &acc,
                      int numLocks, LockAction lockAction,
                      ObjectFifoState &state) {
    ObjectFifoCreateOp target = op;
    auto portNum = port == ObjectFifoPort::Produce ? 0 : 1;
    if (auto linkOp = getOptionalLinkOp(op))
      if (state.objFifoLinks.find(*linkOp) != state.objFifoLinks.end())
        target = state.objFifoLinks[*linkOp];

    auto dev = op->getParentOfType<DeviceOp>();
    if (!dev.getTargetModel().hasProperty(AIETargetModel::UsesSemaphoreLocks)) {

      if (state.locksPerFifo[target].size() == 0) {
        for (int i = 0; i < numLocks; i++) {
          int lockID = acc[{op, portNum}];
          acc[{op, portNum}] =
              (lockID + 1) % op.size(); // update to next objFifo elem
        }
        return;
      }

      int lockMode = 0;
      if ((port == ObjectFifoPort::Produce &&
           lockAction == LockAction::Release) ||
          (port == ObjectFifoPort::Consume &&
           lockAction == LockAction::Acquire))
        lockMode = 1;
      for (int i = 0; i < numLocks; i++) {
        int lockID = acc[{op, portNum}];
        UseLockOp::create(builder, builder.getUnknownLoc(),
                          state.locksPerFifo[target][lockID], lockAction,
                          lockMode);
        acc[{op, portNum}] =
            (lockID + 1) % op.size(); // update to next objFifo elem
      }
    } else {
      if (numLocks == 0)
        return;

      if (state.locksPerFifo[target].size() == 0) {
        acc[{op, portNum}] = (acc[{op, portNum}] + numLocks) %
                             op.size(); // update to next objFifo elem
        return;
      }

      // search for the correct lock based on the port of the acq/rel
      // operation e.g. acq as consumer is the read lock (second)
      LockOp lock;
      if (lockAction == LockAction::AcquireGreaterEqual) {
        if (port == ObjectFifoPort::Produce)
          lock = state.locksPerFifo[target][0];
        else
          lock = state.locksPerFifo[target][1];
      } else {
        if (port == ObjectFifoPort::Produce)
          lock = state.locksPerFifo[target][1];
        else
          lock = state.locksPerFifo[target][0];
      }
      UseLockOp::create(builder, builder.getUnknownLoc(), lock, lockAction,
                        numLocks);
      acc[{op, portNum}] = (acc[{op, portNum}] + numLocks) %
                           op.size(); // update to next objFifo elem
    }
  }

  /// Function used to check whether op is already contained in map.
  /// If it is then return the associated int, if not create new entry and
  /// return 0.
  int updateAndReturnIndex(
      DenseMap<std::pair<ObjectFifoCreateOp, int>, int> &map,
      std::pair<ObjectFifoCreateOp, int> pair) {
    if (map.find(pair) == map.end()) {
      map[pair] = 0;
      return 0;
    }
    return map[pair];
  }

  /// Function used to add an external buffer to the externalBuffersPerFifo map.
  void addExternalBuffer(ObjectFifoCreateOp fifo, ExternalBufferOp buff,
                         ObjectFifoState &state) {
    if (state.externalBuffersPerFifo.find(fifo) ==
        state.externalBuffersPerFifo.end()) {
      std::vector<ExternalBufferOp> buffs;
      state.externalBuffersPerFifo[fifo] = buffs;
    }
    state.externalBuffersPerFifo[fifo].push_back(buff);
  }

  /// Function used to detect all external buffers associated with parent
  /// objectFifo and tile then map them to child objectFifo.
  void detectExternalBuffers(DeviceOp &device, ObjectFifoCreateOp parent,
                             ObjectFifoCreateOp child, Value tile,
                             ObjectFifoState &state) {
    for (auto regOp : device.getOps<ObjectFifoRegisterExternalBuffersOp>())
      if (auto objFifo = regOp.getObjectFifo();
          regOp.getTile() == tile && objFifo == parent)
        for (auto extBuff : regOp.getExternalBuffers())
          addExternalBuffer(child, extBuff.getDefiningOp<ExternalBufferOp>(),
                            state);
  }

  /// Function used to replace uses of split objectFifos.
  void replaceSplitFifo(ObjectFifoCreateOp originalOp, ObjectFifoCreateOp newOp,
                        TileOp tile) {
    auto original =
        originalOp->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName());
    auto newSymbol =
        newOp->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName());
    for (auto user : tile->getUsers())
      if (isa<CoreOp>(user))
        if (auto res =
                SymbolTable::replaceAllSymbolUses(original, newSymbol, user);
            res.failed())
          llvm_unreachable("unreachable");
  }

  /// Function used to find the size of an objectFifo after split based on
  /// the maximum number of elements (of the original objectFifo) acquired
  /// by a process running on given tile. If no CoreOp exists for this tile
  /// return 0.
  int findObjectFifoSize(DeviceOp &device, Value tile,
                         ObjectFifoCreateOp objFifo) {
    if (objFifo.size() == 0)
      return 0;

    // if memTile, size is equal to objFifo size
    if (tile.getDefiningOp<TileOp>().isMemTile())
      return objFifo.size();

    // if shimTile, size is equal to number of external buffers
    if (tile.getDefiningOp<TileOp>().isShimTile())
      for (auto regOp : device.getOps<ObjectFifoRegisterExternalBuffersOp>()) {
        if (regOp.getTile() == tile)
          return regOp.getExternalBuffers().size();
      }

    int maxAcquire = 0;
    for (auto coreOp : device.getOps<CoreOp>())
      if (coreOp.getTile() == tile)
        coreOp.walk([&](ObjectFifoAcquireOp acqOp) {
          if (auto createOp = acqOp.getObjectFifo(); createOp == objFifo)
            if (acqOp.acqNumber() > maxAcquire)
              maxAcquire = acqOp.acqNumber();
        });

    if (maxAcquire > 0) {
      if (maxAcquire == 1 && objFifo.size() == 1)
        return 1;
      return maxAcquire + 1;
      // +1 because objectFifo size is always 1 bigger than maxAcquire to allow
      // for prefetching: simplest case scenario is at least a ping-pong buffer
    }

    return objFifo.size();
  }

  /// Function used to generate, from an objectFifo with a shimTile endpoint, a
  /// shimDMAAllocationOp containing the channelDir, channelIndex and
  /// shimTile reference assigned by the objectFifo lowering.
  void createObjectFifoAllocationInfo(OpBuilder &builder, MLIRContext *ctx,
                                      ObjectFifoCreateOp &objFifoOp,
                                      TileOp shimTile, DMAChannelDir channelDir,
                                      int channelIndex, bool plio,
                                      std::optional<PacketInfoAttr> packet) {
    PacketInfoAttr packetInfo = nullptr;
    if (packet)
      packetInfo = *packet;
    std::string alloc_name = getShimAllocationName(objFifoOp.getName());
    // SymbolRefAttr::get(ctx, objFifoOp.getName())
    ShimDMAAllocationOp::create(
        builder, builder.getUnknownLoc(), StringAttr::get(ctx, alloc_name),
        shimTile.getResult(), DMAChannelDirAttr::get(ctx, channelDir),
        builder.getI64IntegerAttr(channelIndex), builder.getBoolAttr(plio),
        packetInfo);
  }

  static std::string getShimAllocationName(llvm::StringRef objFifoName) {
    return (objFifoName + "_shim_alloc").str();
  }

  /// Function used to verify that an objectfifo is present in at most one
  /// ObjectFifoLinkOp.
  LogicalResult verifyObjectFifoLinks(DeviceOp &device) {
    DenseSet<ObjectFifoCreateOp> objectfifoset;
    bool hasError = false;
    for (ObjectFifoLinkOp link : device.getOps<ObjectFifoLinkOp>()) {
      for (ObjectFifoCreateOp inOf : link.getInputObjectFifos()) {
        if (objectfifoset.count(inOf)) {
          inOf.emitOpError("objectfifo cannot be in more than one "
                           "ObjectFifoLinkOp");
          hasError = true;
        }
        objectfifoset.insert(inOf);
      }
      for (ObjectFifoCreateOp outOf : link.getOutputObjectFifos()) {
        if (objectfifoset.count(outOf)) {
          outOf.emitOpError("objectfifo cannot be in more than one "
                            "ObjectFifoLinkOp");
          hasError = true;
        }
        objectfifoset.insert(outOf);
      }
    }
    return hasError ? failure() : success();
  }

  /// Account for already used packet IDs and return next available ID.
  int getStartPacketID(DeviceOp &device) {
    int packetID = 0;
    for (PacketFlowOp packetflow : device.getOps<PacketFlowOp>()) {
      if (packetflow.getID() > packetID) {
        // compute next available ID
        packetID = packetflow.getID() + 1;
      }
    }
    return packetID;
  }

  /// Helper function to assign DMA channel indices for FIFOs based on
  /// cross-tile conditions
  void assignDMAChannelIndices(
      DMAChannelAnalysis &dmaAnalysis,
      const std::map<ObjectFifoCreateOp, bool> &crossTileInfos,
      std::map<ObjectFifoCreateOp, int> &fifo_dma_channel_index,
      bool assignCrossTileOnly, ObjectFifoState &state) {
    for (auto &[producer, consumers] : state.splitFifos) {
      // Check if we should process this producer based on cross-tile condition
      bool shouldProcessProducer = assignCrossTileOnly
                                       ? crossTileInfos.at(producer)
                                       : !crossTileInfos.at(producer);

      if (shouldProcessProducer) {
        bool requiresAdjacentTileAccessChannels = crossTileInfos.at(producer);
        int channelIndex = dmaAnalysis.getDMAChannelIndex(
            producer.getProducerTileOp(), DMAChannelDir::MM2S,
            requiresAdjacentTileAccessChannels);
        fifo_dma_channel_index[producer] = channelIndex;
      }

      for (auto consumer : consumers) {
        // Check if we should process this consumer based on cross-tile
        // condition
        bool shouldProcessConsumer = assignCrossTileOnly
                                         ? crossTileInfos.at(consumer)
                                         : !crossTileInfos.at(consumer);

        if (shouldProcessConsumer) {
          bool requiresAdjacentTileAccessChannels = crossTileInfos.at(consumer);
          int channelIndex = dmaAnalysis.getDMAChannelIndex(
              consumer.getProducerTileOp(), DMAChannelDir::S2MM,
              requiresAdjacentTileAccessChannels);
          fifo_dma_channel_index[consumer] = channelIndex;
        }
      }
    }
  }

  void runOnOperation() override {

    DeviceOp device = getOperation();

    // Create local state for this device operation - ensures thread and
    // multi-device safety
    ObjectFifoState state;

    LockAnalysis lockAnalysis(device);
    DMAChannelAnalysis dmaAnalysis(device);
    OpBuilder builder = OpBuilder::atBlockTerminator(device.getBody());
    auto ctx = device->getContext();
    auto producerWireType = WireBundle::DMA;
    auto consumerWireType = WireBundle::DMA;
    std::set<TileOp>
        objectFifoTiles; // track cores to check for loops during unrolling

    if (failed(verifyObjectFifoLinks(device)))
      return signalPassFailure();

    //===------------------------------------------------------------------===//
    // Split objectFifos into a consumer end and producer end if needed
    //===------------------------------------------------------------------===//
    // We are going to create additional createObjectFifoOps, so get a copy of
    // all "original" ones before the loop to avoid looping over newly created
    // ones.
    std::vector<ObjectFifoCreateOp> createFifoOps;
    auto range = device.getOps<ObjectFifoCreateOp>();
    createFifoOps.insert(createFifoOps.end(), range.begin(), range.end());
    for (auto createOp : createFifoOps) {
      std::vector<ObjectFifoCreateOp> splitConsumerFifos;
      int consumerIndex = 0;
      int consumerDepth = createOp.size();
      ArrayRef<BDDimLayoutArrayAttr> consumerDims =
          createOp.getDimensionsFromStreamPerConsumer();

      // Only FIFOs using DMA are split into two ends;
      // skip in shared memory case
      if (int share_direction = 0;
          !requiresDMAs(createOp, share_direction, state)) {
        continue;
      }

      for (auto consumerTile : createOp.getConsumerTiles()) {
        auto consumerTileOp = dyn_cast<TileOp>(consumerTile.getDefiningOp());

        if (isa<ArrayAttr>(createOp.getElemNumber())) {
          // +1 to account for 1st depth (producer)
          consumerDepth = createOp.size(consumerIndex + 1);
        } else {
          consumerDepth = findObjectFifoSize(device, consumerTileOp, createOp);
        }

        builder.setInsertionPointAfter(createOp);
        auto datatype = llvm::cast<AIEObjectFifoType>(createOp.getElemType());
        auto consumerObjFifoSize =
            builder.getIntegerAttr(builder.getI32Type(), consumerDepth);
        // rename and replace split objectFifo
        std::string consumerFifoName;
        if (createOp.getConsumerTiles().size() > 1) {
          consumerFifoName = createOp.name().str() + "_" +
                             std::to_string(consumerIndex) + "_cons";
        } else {
          consumerFifoName = createOp.name().str() + "_cons";
        }
        BDDimLayoutArrayAttr emptyDims =
            BDDimLayoutArrayAttr::get(builder.getContext(), {});
        BDDimLayoutArrayAttr singletonFromStreamDims =
            BDDimLayoutArrayAttr::get(
                builder.getContext(),
                ArrayRef<BDDimLayoutAttr>{consumerDims[consumerIndex]});
        BDDimLayoutArrayArrayAttr fromStreamDims =
            BDDimLayoutArrayArrayAttr::get(builder.getContext(),
                                           singletonFromStreamDims);

        ObjectFifoCreateOp consumerFifo = createObjectFifo(
            builder, datatype, consumerFifoName, consumerTile, consumerTile,
            consumerObjFifoSize, emptyDims, fromStreamDims);
        if (createOp.getDisableSynchronization())
          consumerFifo.setDisableSynchronization(true);
        // Propagate iter_count attribute from the original createOp
        // to the new consumerFifo
        if (auto bdChainIterCount = createOp.getIterCount()) {
          consumerFifo.setIterCountAttr(
              builder.getI32IntegerAttr(*bdChainIterCount));
        }
        replaceSplitFifo(createOp, consumerFifo, consumerTileOp);

        // the original createOp to the new consumerFifo. Without this,
        // the consumer-side ObjectFifoCreateOp is attr-less and the
        // downstream propagateSparseCompressionAttr (called from
        // createBdBlock for each consumer-side BD) finds no
        // ``aie.decompress_s2mm`` to read and silently emits no
        // ``aie.enable_compression`` on consumer-side S2MM BDs. The
        // of this propagation; the lit test verified only the BD-emit
        // pass's final hop given a hand-constructed
        // NpuWriteBdOp{aie.enable_compression = true} and never drove
        // the full pipeline through this split-fifo path. See
        // tests/aie2p_microtests/dma_compression_loopback/ for the
        // microtest that surfaced this gap.
        for (StringRef attrName : {"aie.compress_mm2s",
                                   "aie.decompress_s2mm",
                                   "aie.sparsity_pattern",
                                   "aie.sparsity_n",
                                   "aie.sparsity_m",
                                   // VariableRateFifo marker; read by
                                   // unrollForLoops to exclude the fifo
                                   // from LCM-based loop unrolling.
                                   "aie.variable_rate"}) {
          if (auto attr = createOp->getAttr(attrName))
            consumerFifo->setAttr(attrName, attr);
        }

        if (createOp.getAieStream()) {
          int streamEnd = createOp.getAieStream().value();
          if (streamEnd > 0) {
            consumerFifo->setAttr("aie_stream",
                                  builder.getI32IntegerAttr(streamEnd));
            consumerFifo->setAttr(
                "aie_stream_port",
                builder.getI32IntegerAttr(createOp.getAieStreamPort().value()));
          }
          if (streamEnd == 1) {
            createOp->removeAttr("aie_stream");
            createOp->removeAttr("aie_stream_port");
          }
        }

        // identify external buffers that were registered to the consumer fifo
        if (consumerTile.getDefiningOp<TileOp>().isShimTile())
          detectExternalBuffers(device, createOp, consumerFifo, consumerTile,
                                state);

        // record that this objectFifo was split; it will require DMA config
        splitConsumerFifos.push_back(consumerFifo);

        // update the linkOp if the split objFifo was originally its start point
        if (auto linkOp = getOptionalLinkOp(createOp))
          for (ObjectFifoCreateOp fifoIn : linkOp->getInputObjectFifos())
            if (fifoIn.name() == createOp.name() &&
                consumerTile == *linkOp->getOptionalSharedTile())
              if (failed(SymbolTable::replaceAllSymbolUses(
                      createOp, consumerFifo.name(), linkOp->getOperation())))
                llvm::report_fatal_error("unable to update all symbol uses");

        consumerIndex++;
      }

      if (!splitConsumerFifos.empty()) {
        state.splitFifos.emplace_back(createOp, splitConsumerFifos);
      }
    }

    //===------------------------------------------------------------------===//
    // - Create objectFifo buffers and locks.
    // - Populate a list of tiles containing objectFifos for later processing of
    //   the acquires/releases (uses of the FIFO).
    // - Global release counter tracker to keep track of the objectFifo state
    //===------------------------------------------------------------------===//
    for (auto createOp : device.getOps<ObjectFifoCreateOp>()) {

      int share_direction = 0;
      bool shared = !requiresDMAs(createOp, share_direction, state);

      // add all tiles that contain an objectFifo to objectFifoTiles for later
      // loop unrolling pass
      objectFifoTiles.insert(createOp.getProducerTileOp());
      for (auto consumerTile : createOp.getConsumerTiles()) {
        auto consumerTileOp = dyn_cast<TileOp>(consumerTile.getDefiningOp());
        objectFifoTiles.insert(consumerTileOp);
      }

      // identify external buffers that were registered to
      // the producer objectFifo
      if (createOp.getProducerTileOp().isShimTile())
        detectExternalBuffers(device, createOp, createOp,
                              createOp.getProducerTile(), state);

      // if split, the necessary size for producer fifo might change
      if (shared) {
        createObjectFifoElements(builder, lockAnalysis, createOp,
                                 share_direction, state);
      } else {
        if (isa<ArrayAttr>(createOp.getElemNumber()))
          createOp.setElemNumberAttr(
              builder.getI32IntegerAttr(createOp.size()));
        else {
          if (!createOp.getInitValues().has_value()) {

            int prodMaxAcquire = findObjectFifoSize(
                device, createOp.getProducerTileOp(), createOp);
            createOp.setElemNumberAttr(
                builder.getI32IntegerAttr(prodMaxAcquire));
          }
        }
        createObjectFifoElements(builder, lockAnalysis, createOp,
                                 share_direction, state);
      }
    }

    //===------------------------------------------------------------------===//
    // Create flows and tile DMAs
    //===------------------------------------------------------------------===//
    // Only the objectFifos we split above require DMA communication; the others
    // rely on shared memory and share the same buffers.

    // analyze cross-tile buffer allocations and print results
    auto crossTileInfos = analyzeCrossTileFIFOBuffers(state);

    // maps ends of split FIFO to DMA channels
    std::map<ObjectFifoCreateOp, int> fifo_dma_channel_index;

    // assign channel indices for FIFOs with cross-tile issues first
    assignDMAChannelIndices(dmaAnalysis, crossTileInfos, fifo_dma_channel_index,
                            true, state);
    // then assign channel indices for FIFOs without cross-tile issues
    assignDMAChannelIndices(dmaAnalysis, crossTileInfos, fifo_dma_channel_index,
                            false, state);

    int packetID = getStartPacketID(device);
    for (auto &[producer, consumers] : state.splitFifos) {
      int producerChanIndex = -1;
      DMAChannel producerChan;
      PacketFlowOp packetflow;
      if (producer.getAieStream()) {
        int prodStreamEnd = producer.getAieStream().value();
        if (prodStreamEnd == 0 || prodStreamEnd == 2) {
          producerChanIndex = producer.getAieStreamPort().value();
          producerChan = {DMAChannelDir::MM2S, producerChanIndex};
          dmaAnalysis.checkAIEStreamIndex(producer.getProducerTileOp(),
                                          producerChan);
        }
      } else {
        producerChanIndex = fifo_dma_channel_index[producer];
        if (producerChanIndex == -1) {
          producer.getProducerTileOp().emitOpError(
              "number of output DMA channel exceeded!");
          return signalPassFailure();
        }
        producerChan = {DMAChannelDir::MM2S, producerChanIndex};
        std::optional<PacketInfoAttr> bdPacket = {};
        if (clPacketSwObjectFifos) {
          if (packetID > 31) {
            device.emitOpError("max number of packet IDs reached");
            return signalPassFailure();
          }
          bdPacket = {AIE::PacketInfoAttr::get(ctx, /*pkt_type*/ 0,
                                               /*pkt_id*/ packetID)};
          packetID++;
        }
        createDMA(device, builder, producer, producerChan.direction,
                  producerChan.channel, 0, producer.getDimensionsToStreamAttr(),
                  producer.getPadDimensionsAttr(), bdPacket, state);

        // generate objectFifo allocation info
        builder.setInsertionPoint(device.getBody()->getTerminator());
        if (producer.getProducerTileOp().isShimTile())
          createObjectFifoAllocationInfo(
              builder, ctx, producer, producer.getProducerTileOp(),
              producerChan.direction, producerChan.channel, producer.getPlio(),
              bdPacket);

        if (clPacketSwObjectFifos) {
          // create packet flow
          builder.setInsertionPointAfter(producer);
          packetflow = builder.create<PacketFlowOp>(
              builder.getUnknownLoc(),
              builder.getIntegerAttr(builder.getI8Type(), bdPacket->getPktId()),
              nullptr, nullptr);
          {
            OpBuilder::InsertionGuard g(builder);
            builder.setInsertionPointToStart(
                &packetflow.getRegion().emplaceBlock());
            builder.create<EndOp>(builder.getUnknownLoc());
          }
        }
      }

      for (auto consumer : consumers) {
        // if not aie stream, create consumer tile DMA
        int consumerChanIndex = -1;
        DMAChannel consumerChan;
        if (consumer.getAieStream()) {
          int consStreamEnd = consumer.getAieStream().value();
          if (consStreamEnd == 1 || consStreamEnd == 2) {
            consumerChanIndex = consumer.getAieStreamPort().value();
            consumerChan = {DMAChannelDir::S2MM, consumerChanIndex};
            dmaAnalysis.checkAIEStreamIndex(consumer.getProducerTileOp(),
                                            consumerChan);
          }
        } else {
          consumerChanIndex = fifo_dma_channel_index[consumer];
          if (consumerChanIndex == -1) {
            consumer.getProducerTileOp().emitOpError(
                "number of input DMA channel exceeded!");
            return signalPassFailure();
          }
          consumerChan = {DMAChannelDir::S2MM, consumerChanIndex};
          BDDimLayoutArrayAttr consumerDims =
              consumer.getDimensionsFromStreamPerConsumer()[0];
          createDMA(device, builder, consumer, consumerChan.direction,
                    consumerChan.channel, 1, consumerDims, nullptr, {}, state);

          // generate objectFifo allocation info
          builder.setInsertionPoint(device.getBody()->getTerminator());
          if (!consumer.getAieStream()) {
            // generate objectFifo allocation info
            builder.setInsertionPoint(device.getBody()->getTerminator());
            if (consumer.getProducerTileOp().isShimTile())
              createObjectFifoAllocationInfo(
                  builder, ctx, producer, consumer.getProducerTileOp(),
                  consumerChan.direction, consumerChan.channel,
                  producer.getPlio(), {});
          }

          if (clPacketSwObjectFifos) {
            builder.setInsertionPointToStart(&packetflow.getPorts().front());
            builder.create<PacketDestOp>(builder.getUnknownLoc(),
                                         consumer.getProducerTile(),
                                         WireBundle::DMA, consumerChan.channel);
          }
        }

        // If we have PLIO then figure out the direction and make that a PLIO
        if (producer.getPlio()) {
          producerWireType = producer.getProducerTileOp().isShimTile()
                                 ? WireBundle::PLIO
                                 : WireBundle::DMA;
          consumerWireType = consumer.getProducerTileOp().isShimTile()
                                 ? WireBundle::PLIO
                                 : WireBundle::DMA;
        } else {
          producerWireType = WireBundle::DMA;
          consumerWireType = WireBundle::DMA;
          if (producer.getAieStream()) {
            int prodStreamEnd = producer.getAieStream().value();
            if (prodStreamEnd == 0 || prodStreamEnd == 2)
              producerWireType = WireBundle::Core;
          }
          if (consumer.getAieStream()) {
            int consumerStreamEnd = consumer.getAieStream().value();
            if (consumerStreamEnd == 1 || consumerStreamEnd == 2)
              consumerWireType = WireBundle::Core;
          }
        }

        if (!clPacketSwObjectFifos) {
          // create flow
          builder.setInsertionPointAfter(producer);
          FlowOp::create(builder, builder.getUnknownLoc(),
                         producer.getProducerTile(), producerWireType,
                         producerChan.channel, consumer.getProducerTile(),
                         consumerWireType, consumerChan.channel);
        }
      }

      if (clPacketSwObjectFifos) {
        builder.setInsertionPointToStart(&packetflow.getPorts().front());
        PacketSourceOp::create(builder, builder.getUnknownLoc(),
                               producer.getProducerTile(), WireBundle::DMA,
                               producerChan.channel);
      }
    }

    //===------------------------------------------------------------------===//
    // Statically unroll for loops or use dynamic objectFifos
    //===------------------------------------------------------------------===//
    if (clDynamicObjectFifos) {
      if (failed(dynamicGlobalObjectFifos(device, builder, objectFifoTiles,
                                          state)))
        return signalPassFailure();
    } else {
      std::set<TileOp> dynamicTiles;
      std::set<TileOp> unrollTiles;
      for (auto c : device.getOps<CoreOp>()) {
        TileOp t = c.getTileOp();
        if (objectFifoTiles.count(t) > 0) {
          if (c.getDynamicObjfifoLowering().has_value()) {
            if (c.getDynamicObjfifoLowering().value())
              dynamicTiles.insert(t);
            else
              unrollTiles.insert(t);
          } else {
            unrollTiles.insert(t);
          }
        }
      }
      if (failed(
              dynamicGlobalObjectFifos(device, builder, dynamicTiles, state)))
        return signalPassFailure();
      if (failed(unrollForLoops(device, builder, unrollTiles)))
        return signalPassFailure();
    }

    //===------------------------------------------------------------------===//
    // Replace ops
    //===------------------------------------------------------------------===//
    for (auto coreOp : device.getOps<CoreOp>()) {
      DenseMap<ObjectFifoAcquireOp, std::vector<BufferOp *>>
          subviews; // maps each "subview" to its buffer references (subviews
      // are created by AcquireOps)
      DenseMap<std::pair<ObjectFifoCreateOp, int>, std::vector<int>>
          acquiresPerFifo; // maps each objFifo to indices of buffers acquired
      // in latest subview of that objFifo (useful to
      // cascade acquired elements to next AcquireOp)
      DenseMap<std::pair<ObjectFifoCreateOp, int>,
               std::vector<ObjectFifoReleaseOp>>
          releaseOps; // useful to check which ReleaseOp has taken place before
      // an AcquireOp per objFifo
      DenseMap<std::pair<ObjectFifoCreateOp, int>, int>
          acqPerFifo; // maps each objFifo to its next index to acquire within
      // this CoreOp
      DenseMap<std::pair<ObjectFifoCreateOp, int>, int>
          relPerFifo; // maps each objFifo to its next index to release within
      // this CoreOp

      //===----------------------------------------------------------------===//
      // Replace objectFifo.release ops
      //===----------------------------------------------------------------===//
      WalkResult res = coreOp.walk([&](ObjectFifoReleaseOp releaseOp) {
        builder.setInsertionPointAfter(releaseOp);
        ObjectFifoCreateOp op = releaseOp.getObjectFifo();
        auto port = releaseOp.getPort();
        auto portNum = port == ObjectFifoPort::Produce ? 0 : 1;
        auto core = releaseOp->getParentOfType<CoreOp>();

        if (auto linkOp = getOptionalLinkOp(op)) {
          if (core.getTile() == *linkOp->getOptionalSharedTile()) {
            releaseOp->emitOpError("currently cannot access objectFifo used in "
                                   "ObjectFifoLinkOp");
            return WalkResult::interrupt();
            ;
          }
        }

        if (op.getAieStream().has_value()) {
          int streamEnd = op.getAieStream().value();
          if (streamEnd == 2 || streamEnd == portNum)
            releaseOp->emitOpError("cannot release from objectfifo stream "
                                   "port");
          return WalkResult::interrupt();
        }

        // update index of next element to release for this objectFifo
        updateAndReturnIndex(relPerFifo, {op, portNum});

        // release locks
        int numLocks = releaseOp.relNumber();
        // account for repetition
        if (op.getRepeatCount().has_value())
          numLocks *= op.getRepeatCount().value();
        createUseLocks(builder, op, port, relPerFifo, numLocks,
                       LockAction::Release, state);

        // register release op
        if (releaseOps.find({op, portNum}) != releaseOps.end()) {
          releaseOps[{op, portNum}].push_back(releaseOp);
        } else {
          std::vector release = {releaseOp};
          releaseOps[{op, portNum}] = release;
        }
        return WalkResult::advance();
      });
      if (res.wasInterrupted())
        return signalPassFailure();

      //===----------------------------------------------------------------===//
      // Replace objectFifo.acquire ops
      //===----------------------------------------------------------------===//
      res = coreOp.walk([&](ObjectFifoAcquireOp acquireOp) {
        ObjectFifoCreateOp op = acquireOp.getObjectFifo();
        builder.setInsertionPointAfter(acquireOp);
        auto port = acquireOp.getPort();
        auto portNum = port == ObjectFifoPort::Produce ? 0 : 1;
        auto core = acquireOp->getParentOfType<CoreOp>();

        auto linkOp = getOptionalLinkOp(op);
        if (linkOp) {
          if (core.getTile() == *linkOp->getOptionalSharedTile()) {
            acquireOp->emitOpError("currently cannot access objectFifo used in "
                                   "ObjectFifoLinkOp");
            return WalkResult::interrupt();
            ;
          }
        }

        if (op.getAieStream().has_value()) {
          int streamEnd = op.getAieStream().value();
          if (streamEnd == 2 || streamEnd == portNum)
            acquireOp->emitOpError("cannot acquire from objectfifo stream "
                                   "port");
          return WalkResult::interrupt();
        }

        // index of next element to acquire for this objectFifo
        int start = updateAndReturnIndex(
            acqPerFifo, {op, portNum}); // useful for keeping track of which
        // indices are acquired

        // check how many elements have been released in between this AcquireOp
        // and the previous one
        // !!! operations may not be in the same block !!!
        int numRel = 0;
        for (std::vector<ObjectFifoReleaseOp>::iterator relOp =
                 releaseOps[{op, portNum}].begin();
             relOp != releaseOps[{op, portNum}].end();) {
          bool erased = false;
          Operation *acqBlockDefOp = acquireOp.getOperation();
          do {
            Operation *relBlockDefOp = (*relOp).getOperation();
            do {
              if (acqBlockDefOp->getBlock() == relBlockDefOp->getBlock()) {
                if (relBlockDefOp->isBeforeInBlock(acqBlockDefOp)) {
                  numRel += (*relOp).relNumber();
                  relOp = releaseOps[{op, portNum}].erase(relOp);
                  // to ensure that we do not account
                  // the ReleaseOps again later,
                  // after the subview is created
                  erased = true;
                }
              }
            } while ((relBlockDefOp = relBlockDefOp->getParentOp()) &&
                     !isa<DeviceOp>(relBlockDefOp) && !erased);
          } while ((acqBlockDefOp = acqBlockDefOp->getParentOp()) &&
                   !isa<DeviceOp>(acqBlockDefOp) && !erased);
          if (!erased)
            ++relOp;
        }

        // track indices of elements to acquire
        std::vector<int> acquiredIndices;
        if (!acquiresPerFifo[{op, portNum}].empty()) {
          // take into account what has already been acquired by previous
          // AcquireOp in program order
          acquiredIndices = acquiresPerFifo[{op, portNum}];
          // take into account what has been released in-between
          if (static_cast<size_t>(numRel) > acquiredIndices.size()) {
            acquireOp->emitOpError("cannot release more elements than are "
                                   "already acquired");
            return WalkResult::interrupt();
          }
          for (int i = 0; i < numRel; i++)
            acquiredIndices.erase(acquiredIndices.begin());
        }

        // acquire locks
        int numLocks = acquireOp.acqNumber();
        int alreadyAcq = acquiredIndices.size();
        int numCreate;
        if (numLocks > alreadyAcq)
          numCreate = numLocks - alreadyAcq;
        else
          numCreate = 0;

        // account for repetition
        if (op.getRepeatCount().has_value())
          numCreate *= op.getRepeatCount().value();

        auto dev = op->getParentOfType<DeviceOp>();
        if (auto &targetArch = dev.getTargetModel();
            targetArch.getTargetArch() == AIEArch::AIE1)
          createUseLocks(builder, op, port, acqPerFifo, numCreate,
                         LockAction::Acquire, state);
        else
          createUseLocks(builder, op, port, acqPerFifo, numCreate,
                         LockAction::AcquireGreaterEqual, state);

        // if objFifo was linked with others, find which objFifos
        // elements to use
        ObjectFifoCreateOp target = op;
        if (linkOp)
          if (state.objFifoLinks.find(*linkOp) != state.objFifoLinks.end())
            target = state.objFifoLinks[*linkOp];

        // create subview: buffers that were already acquired + new acquires
        for (int i = 0; i < numCreate; i++) {
          acquiredIndices.push_back(start);
          start = (start + 1) % op.size();
        }
        std::vector<BufferOp *> subviewRefs;
        subviewRefs.reserve(acquiredIndices.size());
        for (auto index : acquiredIndices)
          subviewRefs.push_back(&state.buffersPerFifo[target][index]);

        subviews[acquireOp] = subviewRefs;
        acquiresPerFifo[{op, portNum}] = acquiredIndices;

        return WalkResult::advance();
      });
      if (res.wasInterrupted())
        return signalPassFailure();

      //===----------------------------------------------------------------===//
      // Replace subview.access ops
      //===----------------------------------------------------------------===//
      res = coreOp.walk([&](ObjectFifoSubviewAccessOp accessOp) {
        auto acqOp = accessOp.getSubview().getDefiningOp<ObjectFifoAcquireOp>();
        if (ObjectFifoCreateOp op = acqOp.getObjectFifo()) {
          if (auto linkOp = getOptionalLinkOp(op); linkOp.has_value()) {
            if (!linkOp->isDistribute() && !linkOp->isJoin()) {
              for (auto consumerTile : op.getConsumerTiles()) {
                if (auto consumerTileOp =
                        dyn_cast<TileOp>(consumerTile.getDefiningOp())) {
                  int share_dir_value = 0;
                  bool sharing = isSharedMemory(
                      op.getProducerTileOp(), consumerTileOp, &share_dir_value);
                  if (!sharing) {
                    accessOp->emitOpError(
                        "currently cannot access objectFifo used in "
                        "ObjectFifoLinkOp if the tiles don't share memory");
                    return WalkResult::interrupt();
                  }
                }
              }
            } else {
              accessOp->emitOpError(
                  "currently cannot access objectFifo used in "
                  "ObjectFifoLinkOp if it is a distribute or join link");
              return WalkResult::interrupt();
            }
          }
        }
        accessOp.getOutput().replaceAllUsesWith(
            subviews[acqOp][accessOp.getIndex()]->getBuffer());
        return WalkResult::advance();
      });
      if (res.wasInterrupted())
        return signalPassFailure();
    }

    //===------------------------------------------------------------------===//
    // Remove old ops
    //===------------------------------------------------------------------===//
    SetVector<Operation *> opsToErase;
    device.walk([&](Operation *op) {
      if (isa<ObjectFifoLinkOp, ObjectFifoRegisterExternalBuffersOp,
              ObjectFifoAcquireOp, ObjectFifoSubviewAccessOp,
              ObjectFifoReleaseOp, ObjectFifoAllocateOp>(op))
        opsToErase.insert(op);
    });
    SmallVector<Operation *> sorted{opsToErase.begin(), opsToErase.end()};
    computeTopologicalSorting(sorted);
    for (auto *op : llvm::reverse(sorted))
      op->erase();

    //===------------------------------------------------------------------===//
    // Replace any remaining uses of object fifo symbol with symbol of its shim
    // dma allocation.
    //===------------------------------------------------------------------===//
    opsToErase.clear();
    for (auto createOp : device.getOps<ObjectFifoCreateOp>()) {
      std::string shimAllocName = getShimAllocationName(createOp.getName());
      if (failed(SymbolTable::replaceAllSymbolUses(
              createOp.getNameAttr(), builder.getStringAttr(shimAllocName),
              device))) {
        createOp.emitError(
            "failed to replace symbol uses with shim allocation");
        return signalPassFailure();
      }
      opsToErase.insert(createOp);
    }
    for (auto *op : opsToErase) {
      op->erase();
    }
  }
};

std::unique_ptr<OperationPass<DeviceOp>>
AIE::createAIEObjectFifoStatefulTransformPass() {
  return std::make_unique<AIEObjectFifoStatefulTransformPass>();
}