AIEDialect.cpp

//===- AIEDialect.cpp -------------------------------------------*- C++ -*-===//
//
// Copyright (C) 2019-2022 Xilinx, Inc.
// Copyright (C) 2022-2026 Advanced Micro Devices, Inc.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

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

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/DialectImplementation.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Interfaces/FoldInterfaces.h"
#include "mlir/Transforms/InliningUtils.h"

#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/TypeSwitch.h"

using namespace mlir;
using namespace xilinx::AIE;

// Add TableGen'erated dialect definitions (including constructor)
// We implement the initialize() function further below
#include "aie/Dialect/AIE/IR/AIEDialect.cpp.inc"

#define GET_TYPEDEF_CLASSES
#include "aie/Dialect/AIE/IR/AIETypes.cpp.inc"

namespace {

struct AIEInlinerInterface : DialectInlinerInterface {
  using DialectInlinerInterface::DialectInlinerInterface;
  // We don't have any special restrictions on what can be inlined into
  // destination regions. Always allow it.
  bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
                       IRMapping &valueMapping) const final override {
    return true;
  }

  // Operations in aie dialect are always legal to inline since they are
  // pure.
  bool isLegalToInline(Operation *op, Region *, bool wouldBeCloned,
                       IRMapping &) const final override {
    return true;
  }

  // Handle the given inlined terminator by replacing it with a new operation
  // as necessary. Required when the inlined region has more than one block.
  void handleTerminator(Operation *op, Block *newDest) const final override {}

  // Handle the given inlined terminator by replacing it with a new operation
  // as necessary. Required when the region has only one block.
  void handleTerminator(Operation *op,
                        ValueRange valuesToRepl) const final override {}
};

struct AIEDialectFoldInterface : DialectFoldInterface {
  using DialectFoldInterface::DialectFoldInterface;

  /// Registered hook to check if the given region, which is attached to an
  /// operation that is *not* isolated from above, should be used when
  /// materializing constants.
  bool shouldMaterializeInto(Region *region) const final override {
    // Materialize constants into the op that "owns" them rather than letting
    // them hoist up to the enclosing IsolatedFromAbove aie.device:
    //  - aie.core bodies are outlined into standalone funcs, so their
    //    constants must stay local for the func to be self-contained.
    //  - aie.runtime_sequence bodies carry constant operands (e.g. the scalar
    //    fields of npu.* ops). Hoisting them to the device body lets CSE merge
    //    them with a core's identical constants, which would leave the core
    //    referencing a device-level value that is erased when the core is
    //    outlined.
    //  - Make sure SSA values for aie.use_lock operands in
    //    aie.mem/aie.memtile_dma/aie.shim_dma bodies do not get
    //    hoisted.
    return isa<CoreOp, RuntimeSequenceOp, MemOp, MemTileDMAOp, ShimDMAOp>(
        region->getParentOp());
  }
};

} // end anonymous namespace

void AIEDialect::initialize() {
  addTypes<
#define GET_TYPEDEF_LIST
#include "aie/Dialect/AIE/IR/AIETypes.cpp.inc"
      >();
  addAttributes<
#define GET_ATTRDEF_LIST
#include "aie/Dialect/AIE/IR/AIEAttrs.cpp.inc"
      >();
  addOperations<
#define GET_OP_LIST
#include "aie/Dialect/AIE/IR/AIEOps.cpp.inc"
      >();
  addInterfaces<AIEInlinerInterface, AIEDialectFoldInterface>();
}

// Helper methods to retrieve the encoding associated to a burst length,
// or to find the highest available burst length if the requested one is 0
// (default value).

static std::pair<uint32_t, uint32_t>
getShimBurstLength(const xilinx::AIE::AIETargetModel &tm,
                   uint32_t burstLength) {

  std::vector<std::pair<uint32_t, uint32_t>> bel =
      tm.getShimBurstEncodingsAndLengths();

  // If we have the default burst length (no burst length was specified),
  // use the highest one available on our target model
  if (burstLength == 0) {
    return *std::max_element(
        bel.begin(), bel.end(),
        [](auto pair1, auto pair2) { return pair1.second < pair2.second; });
  }

  // Note that if we are given a burst size, we are checking its existence in
  // the pass verification already, so we can safely assume it exists.
  return *std::find_if(bel.begin(), bel.end(),
                       [=](auto p) { return p.second == burstLength; });
}

uint32_t xilinx::AIE::getShimBurstLengthBytes(const AIE::AIETargetModel &tm,
                                              uint32_t burstLength) {

  return getShimBurstLength(tm, burstLength).second;
}

uint32_t xilinx::AIE::getShimBurstLengthEncoding(const AIE::AIETargetModel &tm,
                                                 uint32_t burstLength) {

  return getShimBurstLength(tm, burstLength).first;
}

std::string xilinx::AIE::generateUniqueSymbolName(
    mlir::Operation *symbolTableOp, llvm::StringRef prefix, unsigned &counter) {
  std::string name;
  do {
    name = (prefix + llvm::Twine(counter++)).str();
  } while (mlir::SymbolTable::lookupSymbolIn(symbolTableOp, name));
  return name;
}

LogicalResult
xilinx::AIE::myVerifyOffsetSizeAndStrideOp(OffsetSizeAndStrideOpInterface op) {
  std::array<unsigned, 3> maxRanks = op.getArrayAttrMaxRanks();
  if (!(op.getMixedOffsets().size() == 1 && maxRanks[0] == 1) && // NOLINT
      op.getMixedOffsets().size() != op.getMixedSizes().size())
    return op->emitError(
               "expected mixed offsets rank to match mixed sizes rank (")
           << op.getMixedOffsets().size() << " vs " << op.getMixedSizes().size()
           << ") so the rank of the result type is well-formed.";
  if (failed(verifyListOfOperandsOrIntegers(
          op, "offset", maxRanks[0], op.getStaticOffsets(), op.getOffsets())))
    return failure();
  if (failed(verifyListOfOperandsOrIntegers(
          op, "size", maxRanks[1], op.getStaticSizes(), op.getSizes())))
    return failure();
  if (failed(verifyListOfOperandsOrIntegers(
          op, "stride", maxRanks[2], op.getStaticStrides(), op.getStrides())))
    return failure();
  for (int64_t offset : op.getStaticOffsets())
    if (offset < 0 && !ShapedType::isDynamic(offset))
      return op->emitError("expected offsets to be non-negative, but got ")
             << offset;
  for (int64_t size : op.getStaticSizes())
    if (size < 0 && !ShapedType::isDynamic(size))
      return op->emitError("expected sizes to be non-negative, but got ")
             << size;

  return success();
}

static VC1902TargetModel VC1902model;
static VE2302TargetModel VE2302model;
static VE2802TargetModel VE2802model;
static VirtualizedNPU1TargetModel NPUmodel1col(1);
static VirtualizedNPU1TargetModel NPUmodel2col(2);
static VirtualizedNPU1TargetModel NPUmodel3col(3);
static VirtualizedNPU1TargetModel NPUmodel4col(4);
static NPU2TargetModel NPU2model;
static VirtualizedNPU2TargetModel NPU2model1col(1);
static VirtualizedNPU2TargetModel NPU2model2col(2);
static VirtualizedNPU2TargetModel NPU2model3col(3);
static VirtualizedNPU2TargetModel NPU2model4col(4);
static VirtualizedNPU2TargetModel NPU2model5col(5);
static VirtualizedNPU2TargetModel NPU2model6col(6);
static VirtualizedNPU2TargetModel NPU2model7col(7);

const AIETargetModel &xilinx::AIE::getTargetModel(Operation *op) {
  if (auto t = dyn_cast<AIETarget>(op))
    return t.getTargetModel();
  if (auto t = op->getParentOfType<AIETarget>())
    return t.getTargetModel();

  // For backward compatibility, return a basic device model compatible with
  // the VCK190
  return VC1902model;
}

const AIETargetModel &xilinx::AIE::getTargetModel(AIEDevice device) {
  switch (device) {
  case AIEDevice::xcvc1902:
    return VC1902model;
  case AIEDevice::xcve2302:
    return VE2302model;
  case AIEDevice::xcve2802:
    return VE2802model;
  case AIEDevice::npu1:
    return NPUmodel4col;
  case AIEDevice::npu1_1col:
    return NPUmodel1col;
  case AIEDevice::npu1_2col:
    return NPUmodel2col;
  case AIEDevice::npu1_3col:
    return NPUmodel3col;
  case AIEDevice::npu2:
    return NPU2model;
  case AIEDevice::npu2_1col:
    return NPU2model1col;
  case AIEDevice::npu2_2col:
    return NPU2model2col;
  case AIEDevice::npu2_3col:
    return NPU2model3col;
  case AIEDevice::npu2_4col:
    return NPU2model4col;
  case AIEDevice::npu2_5col:
    return NPU2model5col;
  case AIEDevice::npu2_6col:
    return NPU2model6col;
  case AIEDevice::npu2_7col:
    return NPU2model7col;
  }
  return VC1902model;
}

// Walk the operation hierarchy until we find a containing TileElement.
// If no parent is a TileElement, then return null.
static TileElement getParentTileElement(Operation *op) {
  auto *parent = op->getParentOp();
  while (!llvm::isa_and_nonnull<DeviceOp, ModuleOp>(parent)) {
    if (auto element = llvm::dyn_cast<TileElement>(parent))
      return element;
    parent = parent->getParentOp();
  }
  return llvm::dyn_cast<TileElement>(parent);
}

// Returns the maximum index described by the input dimensions.
static int64_t getDimsMaxIdx(ArrayRef<BDDimLayoutAttr> dims) {
  int64_t maxIdx = 0;
  for (BDDimLayoutAttr dim : dims) {
    maxIdx += dim.getStride() * (dim.getSize() - 1);
  }
  return maxIdx;
}

namespace {

struct UsesAreAccessible {
  static LogicalResult verifyTrait(Operation *op) {
    auto thisElement = cast<TileElement>(op);

    // Skip accessibility checks for logical tiles as we cannot tell until tile
    // is placed
    if (!isa<TileOp>(thisElement.getTile().getDefiningOp()))
      return success();

    auto thisID = thisElement.getTileID();
    auto users = op->getResult(0).getUsers();
    const auto &targetModel = getTargetModel(op);
    for (auto *user : users) {
      // AIE.useLock may be used in a device to set the lock's default value
      // Allow in a toplevel module for backward compatibility
      if (llvm::isa_and_nonnull<DeviceOp, ModuleOp>(user->getParentOp())) {
        continue;
      }
      // If any parent or the user itself prescribe that accessibility checks be
      // skipped, skip the check for that user.
      if (user->getParentWithTrait<SkipAccessibilityCheckTrait>() ||
          user->hasTrait<SkipAccessibilityCheckTrait>()) {
        continue;
      }
      TileElement element = llvm::dyn_cast<TileElement>(user);
      if (!element) {
        element = getParentTileElement(user);
      }
      if (!element) {
        // This should probably be caught elsewhere as well.
        return op->emitOpError("is accessed outside of a tile")
                   .attachNote(user->getLoc())
               << "user";
      }
      auto tileID = element.getTileID();
      if (!targetModel.isLegalMemAffinity(tileID.col, tileID.row, thisID.col,
                                          thisID.row)) {
        return (op->emitOpError("in Column ")
                << thisID.col << " and Row " << thisID.row
                << " is accessed from an unreachable tile in Column "
                << tileID.col << " and Row " << tileID.row)
                   .attachNote(user->getLoc())
               << "user";
      }
    }
    return success();
  }
};

} // namespace

// Check that the operation only contains terminators in
// TerminatorOpTypes.
template <typename... TerminatorOpTypes>
struct HasSomeTerminator {
  static LogicalResult verifyTrait(Operation *op) {
    for (auto &region : op->getRegions()) {
      for (auto &block : region) {
        if (!block.empty()) {
          if (Operation *operation = &block.back();
              !llvm::isa_and_nonnull<TerminatorOpTypes...>(operation))
            return operation->emitOpError("is not an allowed terminator")
                .attachNote(op->getLoc())
                .append("in this context: ");
        }
      }
    }
    return success();
  }
};

// Check that the given DMA-like op (e.g. MemOp, ShimDMAOp)
// has valid BDs.
template <typename ConcreteType>
LogicalResult HasValidBDs<ConcreteType>::verifyTrait(Operation *op) {
  auto element = cast<ConcreteType>(op);
  const auto &targetModel = getTargetModel(op);

  TileLike tile = element.getTileLike();
  if (!tile)
    return op->emitOpError("tile must implement TileLike interface");

  int bdMax = targetModel.getNumBDs(tile.getTileType());

  int bdNum = 0;
  for (auto &block : element.getBody()) {
    auto bdOps = llvm::to_vector_of<DMABDOp>(block.template getOps<DMABDOp>());

    // Skip entry/end block
    if (bdOps.empty())
      continue;

    // Check BD count limit
    if (bdNum >= bdMax) {
      return (op->emitOpError("has more than ") << bdMax << " blocks")
          .attachNote(bdOps.front().getLoc())
          .append("no space for this BD");
    }
    bdNum++;

    // Check exactly 1 DMABDOp per BD block
    if (bdOps.size() != 1) {
      return (op->emitOpError("BD block must have exactly one DMABDOp, found ")
              << bdOps.size())
          .attachNote(block.front().getLoc())
          .append("in this BD block");
    }

    // Check at most 2 UseLockOps per BD block (1 acquire, 1 release)
    auto useLockOps =
        llvm::to_vector_of<UseLockOp>(block.template getOps<UseLockOp>());
    int acquireCount = 0;
    int releaseCount = 0;
    for (auto useLock : useLockOps) {
      if (useLock.acquire() || useLock.acquireGE())
        acquireCount++;
      else if (useLock.release())
        releaseCount++;
    }

    if (acquireCount > 1) {
      return (op->emitOpError(
                  "BD block must have at most one acquire UseLockOp, found ")
              << acquireCount)
          .attachNote(block.front().getLoc())
          .append("in this BD block");
    }
    if (releaseCount > 1) {
      return (op->emitOpError(
                  "BD block must have at most one release UseLockOp, found ")
              << releaseCount)
          .attachNote(block.front().getLoc())
          .append("in this BD block");
    }
  }
  return success();
}

// Check that the given DMA-like op (e.g. MemOp, ShimDMAOp)
// has valid DMA channels.
template <typename ConcreteType>
LogicalResult HasValidDMAChannels<ConcreteType>::verifyTrait(Operation *op) {
  DenseSet<DMAChannel> inputChannels;
  DenseSet<DMAChannel> outputChannels;
  auto element = cast<ConcreteType>(op);
  Region &body = element.getBody();
  if (body.empty())
    return op->emitOpError("should have non-empty body");
  for (auto &bodyOp : body.getOps()) {
    // check for duplicate DMA channels within the same ShimDMAOp
    if (auto dmaStart = dyn_cast<DMAStartOp>(bodyOp)) {
      DMAChannel dmaChan = {dmaStart.getChannelDir(),
                            dmaStart.getChannelIndex()};
      // check if number of input and output channels is more than available
      // hardware
      if (dmaChan.direction == DMAChannelDir::S2MM)
        inputChannels.insert(dmaChan);
      else
        outputChannels.insert(dmaChan);
    }
  }

  TileLike tile = element.getTileLike();
  if (!tile)
    return op->emitOpError("tile must implement TileLike interface");

  if (inputChannels.size() > tile.getNumSourceConnections(WireBundle::DMA))
    return op->emitOpError(
        "uses more input channels than available on this tile");

  if (outputChannels.size() > tile.getNumDestConnections(WireBundle::DMA))
    return op->emitOpError(
        "uses more output channels than available on this tile");
  return success();
}

//===----------------------------------------------------------------------===//
// ObjectFifoCreateOp
//===----------------------------------------------------------------------===//

LogicalResult ObjectFifoCreateOp::verify() {
  if (isa<ArrayAttr>(getElemNumber())) {
    if (size_t numDepths = dyn_cast<ArrayAttr>(getElemNumber()).size();
        numDepths != getConsumerTiles().size() + 1) // +1 for producer depth
      return emitOpError("does not have enough depths specified for producer "
                         "and for each consumer.");
  }

  // Helper to get tile interface from Value
  auto getTileLikeFromValue = [](Value v) -> TileLike {
    return llvm::dyn_cast<TileLike>(v.getDefiningOp());
  };

  TileLike producerTile = getTileLikeFromValue(getProducerTile());
  if (!producerTile)
    return emitError("producer tile must implement TileLike interface");

  // data layout transformations on shim tiles are handled by runtime operations
  if (producerTile.isShimTile() && !getDimensionsToStream().empty()) {
    return emitError(
        "`dimensionsToStream` data layout transformations are not supported "
        "on shim tile producers");
  }
  for (auto consTileVal : getConsumerTiles()) {
    TileLike consTile = getTileLikeFromValue(consTileVal);
    if (!consTile)
      return emitError("consumer tile must implement TileLike interface");
    if (consTile.isShimTile() &&
        !getDimensionsFromStream(consTileVal).empty()) {
      return emitError(
          "`dimensionsFromStreamPerConsumer` data layout transformations are "
          "not supported on shim tile consumers");
    }
  }

  if (getRepeatCount().has_value()) {
    if (producerTile.isShimTile())
      return emitError("`repeat_count` unavailable for shim tiles");
  }

  if (getAieStreamPort().has_value()) {
    if (!getAieStream().has_value())
      return emitError("`aie_stream` must be defined");
  }

  if (getAieStream().has_value()) {
    if (getConsumerTiles().size() > 1)
      return emitError("`aie_stream` can only be used in 1-to-1 object FIFOs");

    if (!getAieStreamPort().has_value())
      return emitError("`aie_stream_port` must be defined");

    if (getAieStream().value() == 0 || getAieStream().value() == 2) {
      if (producerTile.isShimTile() || producerTile.isMemTile())
        return emitError(
            "`aie_stream` is not available for shim and mem tiles");

      if (getRepeatCount().has_value())
        return emitError("`repeat_count` unavailable on stream end");

      if (getInitValues().has_value())
        return emitError("`init_values` unavailable on stream end");

      if (getIterCount().has_value())
        return emitError("`iter_count` unavailable on stream end");

      if (!getDimensionsToStream().empty())
        return emitError("`dimensionsToStream` data layout transformations are "
                         "unavailable on stream end");
    }

    if (getAieStream().value() == 1 || getAieStream().value() == 2) {
      TileLike consTile = getTileLikeFromValue(getConsumerTiles()[0]);
      if (consTile && (consTile.isShimTile() || consTile.isMemTile()))
        return emitError(
            "`aie_stream` is not available for shim and mem tiles");
    }

    if (!getDimensionsFromStreamPerConsumer()[0].empty())
      return emitError("`dimensionsFromStreamPerConsumer` data layout "
                       "transformations are unavailable on stream end");
  }

  if (getInitValues().has_value()) {
    if (producerTile.isShimTile())
      return emitError("`init_values` unavailable for shim tiles");
  }

  if (getInitValues().has_value()) {
    if ((int)getInitValues().value().size() != size())
      return emitError("`init_values` does not initialize all objects");
  }

  if (getIterCount().has_value()) {
    int iterCount = getIterCount().value();
    if (iterCount < 1 || iterCount > 256)
      return emitError("`iter_count` must be between 1 and 256");

    // Check that either producer or at least one consumer is a MemTile
    bool hasMemTile = producerTile.isMemTile();
    if (!hasMemTile) {
      for (auto consTileVal : getConsumerTiles()) {
        TileLike consTile = getTileLikeFromValue(consTileVal);
        if (consTile && consTile.isMemTile()) {
          hasMemTile = true;
          break;
        }
      }
    }
    if (!hasMemTile)
      return emitError("`iter_count` is currently only supported on MemTiles");
  }

  if (getConsumerElemType().has_value()) {
    auto consType =
        llvm::dyn_cast<AIEObjectFifoType>(getConsumerElemType().value());
    if (!consType)
      return emitError("consumer element type must be an "
                       "!aie.objectfifo<memref<...>> type");
    auto prodType = llvm::cast<AIEObjectFifoType>(getElemType());
    auto prodMemref = prodType.getElementType();
    auto consMemref = consType.getElementType();
    if (prodMemref.getElementType() != consMemref.getElementType())
      return emitError("producer and consumer must have the same scalar "
                       "element type, but got ")
             << prodMemref.getElementType() << " vs "
             << consMemref.getElementType();
    int64_t prodSize = prodMemref.getNumElements();
    int64_t consSize = consMemref.getNumElements();
    if (consSize <= 0)
      return emitError("consumer element count must be positive");
    if (prodSize % consSize != 0)
      return emitError("producer element size (")
             << prodSize << ") must be an integer multiple of consumer "
             << "element size (" << consSize << ")";
  }

  return success();
}

TileOp ObjectFifoCreateOp::getProducerTileOp() {
  return cast<TileOp>(getProducerTile().getDefiningOp());
}

BDDimLayoutArrayAttr
ObjectFifoCreateOp::getDimensionsFromStream(Value consumerTile) {
  int dimsIndex = 0;
  for (auto cons : getConsumerTiles()) {
    if (cons == consumerTile)
      break;
    else
      dimsIndex++;
  }
  return getDimensionsFromStreamPerConsumer()[dimsIndex];
}

ParseResult xilinx::AIE::parseObjectFifoProducerTile(
    OpAsmParser &parser, OpAsmParser::UnresolvedOperand &operand,
    BDDimLayoutArrayAttr &dimensions) {
  std::vector<BDDimLayoutAttr> emptyDims = {};
  if (parser.parseOperand(operand))
    return failure();
  if (succeeded(parser.parseOptionalKeyword("dimensionsToStream"))) {
    if (parser.parseCustomAttributeWithFallback<BDDimLayoutArrayAttr>(
            dimensions)) {
      return failure();
    }
  } else {
    dimensions =
        BDDimLayoutArrayAttr::get(parser.getContext(), ArrayRef(emptyDims));
  }
  return success();
}

void xilinx::AIE::printObjectFifoProducerTile(OpAsmPrinter &printer,
                                              Operation *op, Value operand,
                                              BDDimLayoutArrayAttr dimensions) {
  printer << operand;
  if (!dimensions.empty()) {
    printer << " dimensionsToStream ";
    printer.printStrippedAttrOrType(dimensions);
  }
}

ParseResult xilinx::AIE::parseObjectFifoConsumerTiles(
    OpAsmParser &parser, SmallVectorImpl<OpAsmParser::UnresolvedOperand> &tiles,
    BDDimLayoutArrayArrayAttr &dimensions) {
  // parseCommaSeparatedList doesn't handle the missing case for "none",
  // so we handle it custom here.
  std::vector<BDDimLayoutArrayAttr> tileDims = {};

  auto parseOneOperand = [&]() -> ParseResult {
    if (parser.parseOperand(tiles.emplace_back(), true)) {
      return failure();
    }
    // By default, create empty dimensions array for each consumer; this way,
    // we can be certain to have as many entries in the dimensions array as
    // there are customer
    BDDimLayoutArrayAttr dimAttr =
        BDDimLayoutArrayAttr::get(parser.getContext(), {});

    if (succeeded(parser.parseOptionalKeyword("dimensionsFromStream"))) {
      // If specified, parse actual data layout transform dimensions
      if (parser.parseCustomAttributeWithFallback<BDDimLayoutArrayAttr>(
              dimAttr)) {
        return failure();
      }
    }
    tileDims.emplace_back(dimAttr);
    return success();
  };

  if (parser.parseCommaSeparatedList(AsmParser::Delimiter::None,
                                     parseOneOperand, " in operand list"))
    return failure();

  dimensions = BDDimLayoutArrayArrayAttr::get(parser.getContext(), tileDims);
  return success();
}

void xilinx::AIE::printObjectFifoConsumerTiles(
    OpAsmPrinter &printer, Operation *op, OperandRange tiles,
    BDDimLayoutArrayArrayAttr dimsPerTileAttr) {
  size_t tileIdx = 0;
  for (auto tile : tiles) {
    printer << tile;
    if (dimsPerTileAttr && tileIdx < dimsPerTileAttr.size() &&
        dimsPerTileAttr[tileIdx] && !dimsPerTileAttr[tileIdx].empty()) {
      printer << " dimensionsFromStream ";
      printer.printStrippedAttrOrType(dimsPerTileAttr[tileIdx]);
    }
    if (tileIdx < tiles.size() - 1) {
      printer << ", ";
    }
    tileIdx++;
  }
}

static void printObjectFifoConsumerElemType(OpAsmPrinter &p,
                                            ObjectFifoCreateOp op,
                                            TypeAttr consumerElemType) {
  if (consumerElemType)
    p << " -> " << consumerElemType;
}

static ParseResult parseObjectFifoConsumerElemType(OpAsmParser &parser,
                                                   TypeAttr &consumerElemType) {
  if (failed(parser.parseOptionalArrow()))
    return success(); // no consumer type
  Type type;
  if (parser.parseType(type))
    return failure();
  consumerElemType = TypeAttr::get(type);
  return success();
}

static void printObjectFifoInitValues(OpAsmPrinter &p, ObjectFifoCreateOp op,
                                      Attribute numElem, TypeAttr type,
                                      Attribute initValues) {
  if (op.getInitValues()) {
    p << "= [";
    // `numElem` may be an IntegerAttr (scalar depth) or an ArrayAttr of
    // per-endpoint depths; initValues populates the producer side, which
    // is the first entry of the ArrayAttr.
    int depth;
    if (isa<ArrayAttr>(numElem)) {
      depth = llvm::dyn_cast<mlir::IntegerAttr>(
                  llvm::dyn_cast<mlir::ArrayAttr>(numElem)[0])
                  .getInt();
    } else {
      depth = llvm::dyn_cast<mlir::IntegerAttr>(numElem).getInt();
    }
    for (int i = 0; i < depth; i++) {
      p.printStrippedAttrOrType(llvm::dyn_cast<mlir::ArrayAttr>(initValues)[i]);
      if (i < depth - 1) {
        p << ", ";
      }
    }
    p << "]";
  }
}

static ParseResult parseObjectFifoInitValues(OpAsmParser &parser,
                                             Attribute numElem, TypeAttr type,
                                             Attribute &initValues) {
  int depth;
  if (isa<ArrayAttr>(numElem)) {
    depth = llvm::dyn_cast<mlir::IntegerAttr>(
                llvm::dyn_cast<mlir::ArrayAttr>(numElem)[0])
                .getInt();
  } else {
    depth = llvm::dyn_cast<mlir::IntegerAttr>(numElem).getInt();
  }
  auto objfifoType = llvm::cast<AIEObjectFifoType>(type.getValue());
  auto memrefType = llvm::cast<MemRefType>(objfifoType.getElementType());

  if (!memrefType.hasStaticShape())
    return parser.emitError(parser.getNameLoc())
           << "type should be static shaped memref, but got " << memrefType;

  if (parser.parseOptionalEqual())
    return success();

  Type tensorType = mlir::memref::getTensorTypeFromMemRefType(memrefType);
  if (parser.parseAttribute(initValues, tensorType))
    return failure();
  for (int i = 0; i < depth; i++) {
    auto initialValues = llvm::dyn_cast<mlir::ArrayAttr>(initValues);
    if ((int)initialValues.size() != depth)
      return parser.emitError(parser.getNameLoc())
             << "initial values should initialize all objects";
    if (!llvm::isa<ElementsAttr>(initialValues[i]))
      return parser.emitError(parser.getNameLoc())
             << "initial value should be an elements attribute";
  }

  return success();
}

//===----------------------------------------------------------------------===//
// ObjectFifoAllocateOp
//===----------------------------------------------------------------------===//

LogicalResult ObjectFifoAllocateOp::verify() {
  ObjectFifoCreateOp objFifo = getObjectFifo();
  if (!objFifo)
    return emitError("cannot retrieve associated object FIFO");
  if (objFifo.getConsumerTiles().size() != 1)
    return emitError("can only be used in 1-to-1 object FIFOs");
  if (objFifo.getVia_DMA())
    return emitError("cannot allocate a shared memory module to objectfifo "
                     "with set `via_DMA` attribute");
  if (objFifo.getRepeatCount().has_value())
    return emitError("cannot allocate a shared memory module to objectfifo "
                     "with set `repeat_count` attribute");
  if (!objFifo.getDimensionsToStream().empty())
    return emitError("cannot allocate a shared memory module to objectfifo "
                     "with set dimensions attribute");
  if (objFifo.getAieStream().has_value())
    return emitError("cannot allocate a shared memory module to objectfifo "
                     "using stream port");
  return success();
}

TileOp ObjectFifoAllocateOp::getDelegateTileOp() {
  return cast<TileOp>(getDelegateTile().getDefiningOp());
}

ObjectFifoCreateOp ObjectFifoAllocateOp::getObjectFifo() {
  Operation *parent = getOperation();
  while ((parent = parent->getParentOp())) {
    if (parent->hasTrait<OpTrait::SymbolTable>()) {
      if (auto *st = SymbolTable::lookupSymbolIn(parent, getObjFifoName());
          isa_and_nonnull<ObjectFifoCreateOp>(st))
        return dyn_cast<ObjectFifoCreateOp>(st);
    }
  }
  return {};
}

//===----------------------------------------------------------------------===//
// ObjectFifoLinkOp
//===----------------------------------------------------------------------===//

LogicalResult ObjectFifoLinkOp::verify() {
  if (isJoin() && isDistribute())
    return emitError("ObjectFifoLinkOp does not support 'join' and "
                     "'distribute' at the same time");

  auto sharedTile = getOptionalSharedTile();
  if (!sharedTile)
    return emitError("ObjectFifoLinkOp must have a link point, i.e., a "
                     "shared tile between objectFifos");

  TileLike tile = llvm::dyn_cast<TileLike>(sharedTile.value().getDefiningOp());
  if (!tile)
    return emitError("shared tile must implement TileLike interface");
  if (!tile.isMemTile()) {
    if (isJoin() || isDistribute())
      return emitError("ObjectFifoLinkOp join and distribute are "
                       "unavailable on compute or shim tiles");
  }

  if (isJoin()) {
    if (getFifoIns().size() != getSrcOffsets().size())
      return emitOpError("number of provided src offsets must be equal "
                         "to the number of input objectFifos");

    if (!getDstOffsets().empty())
      return emitOpError("dst offsets should be empty for join");

    ObjectFifoCreateOp fifoOut = getOutputObjectFifos()[0];
    if (!fifoOut.getDimensionsToStream().empty()) {
      int64_t maxIdx = getDimsMaxIdx(fifoOut.getDimensionsToStream());
      int64_t minInputBufferSize = -1;
      for (auto lenIn : getJoinTransferLengths()) {
        if (lenIn <= minInputBufferSize || minInputBufferSize < 0)
          minInputBufferSize = lenIn;
      }
      if (minInputBufferSize <= maxIdx) {
        return emitOpError()
               << "specified output stride(s) and size(s) result in out "
                  "of bounds access in join input, for index "
               << std::to_string(maxIdx) << " in transfer of length "
               << std::to_string(minInputBufferSize) << ".";
      }
    }

  } else if (isDistribute()) {
    if (getFifoOuts().size() != getDstOffsets().size())
      return emitOpError("number of provided dst offsets must be equal "
                         "to the number of output objectFifos");

    if (!getSrcOffsets().empty())
      return emitOpError("src offsets should be empty for distribute");

    ObjectFifoCreateOp fifoIn = getInputObjectFifos()[0];
    if (!fifoIn.getDimensionsFromStream(sharedTile.value()).empty()) {
      int64_t maxIdx =
          getDimsMaxIdx(fifoIn.getDimensionsFromStream(sharedTile.value()));
      int64_t minOutputBufferSize = -1;
      for (auto lenOut : getDistributeTransferLengths()) {
        if (lenOut <= minOutputBufferSize || minOutputBufferSize < 0)
          minOutputBufferSize = lenOut;
      }
      if (minOutputBufferSize <= maxIdx) {
        return emitOpError()
               << "specified input stride(s) and size(s) result in out "
                  "of bounds access in distribute output, for index "
               << std::to_string(maxIdx) << " in transfer of length "
               << std::to_string(minOutputBufferSize) << ".";
      }
    }

    std::vector<int> repeat_counts;
    for (auto fifoOut : getOutputObjectFifos()) {
      if (fifoOut.getRepeatCount().has_value()) {
        repeat_counts.push_back(fifoOut.getRepeatCount().value());
      } else {
        repeat_counts.push_back(0);
      }
    }
    for (auto repeat : repeat_counts)
      if (repeat_counts[0] != repeat)
        return emitError("repeat counts of output object FIFOs must be equal");

  } else {
    if (!getSrcOffsets().empty() && !getDstOffsets().empty())
      return emitOpError("all offsets should be empty if there is no "
                         "join or distribute");
  }

  return success();
}

std::optional<Value> ObjectFifoLinkOp::getOptionalSharedTile() {
  if (isJoin()) {
    auto fifoOut = getOutputObjectFifos()[0];
    for (auto fifoIn : getInputObjectFifos())
      if (fifoOut.getProducerTile() != fifoIn.getConsumerTiles()[0])
        return {};
    return {fifoOut.getProducerTile()};
  }

  if (isDistribute()) {
    auto fifoIn = getInputObjectFifos()[0];
    for (auto fifoOut : getOutputObjectFifos())
      if (fifoIn.getConsumerTiles()[0] != fifoOut.getProducerTile())
        return {};
    return {fifoIn.getConsumerTiles()[0]};
  }

  auto fifoIn = getInputObjectFifos();
  if (auto fifoOut = getOutputObjectFifos();
      !fifoIn.empty() && !fifoOut.empty())
    for (auto consumerIn : fifoIn[0].getConsumerTiles())
      if (consumerIn == fifoOut[0].getProducerTile())
        return {fifoOut[0].getProducerTile()};
  return {};
}

std::vector<ObjectFifoCreateOp> ObjectFifoLinkOp::getInputObjectFifos() {
  std::vector<ObjectFifoCreateOp> inputObjFifos;
  Operation *parent = getOperation();
  while ((parent = parent->getParentOp())) {
    if (parent->hasTrait<OpTrait::SymbolTable>()) {
      for (auto sym : getFifoIns()) {
        auto name = dyn_cast<FlatSymbolRefAttr>(sym);
        if (auto *st = SymbolTable::lookupSymbolIn(parent, name);
            isa_and_nonnull<ObjectFifoCreateOp>(st))
          inputObjFifos.push_back(dyn_cast<ObjectFifoCreateOp>(st));
      }
    }
  }
  return inputObjFifos;
}

std::vector<ObjectFifoCreateOp> ObjectFifoLinkOp::getOutputObjectFifos() {
  std::vector<ObjectFifoCreateOp> outputObjFifos;
  Operation *parent = getOperation();
  while ((parent = parent->getParentOp())) {
    if (parent->hasTrait<OpTrait::SymbolTable>()) {
      for (auto sym : getFifoOuts()) {
        auto name = dyn_cast<FlatSymbolRefAttr>(sym);
        if (auto *st = mlir::SymbolTable::lookupSymbolIn(parent, name);
            isa_and_nonnull<ObjectFifoCreateOp>(st))
          outputObjFifos.push_back(dyn_cast<ObjectFifoCreateOp>(st));
      }
    }
  }
  return outputObjFifos;
}

std::vector<int> ObjectFifoLinkOp::getJoinTransferLengths() {
  std::vector<int> lengths;
  if (isJoin()) {
    auto fifoOut =
        llvm::cast<AIEObjectFifoType>(getOutputObjectFifos()[0].getElemType());
    auto elemTypeOut = llvm::cast<MemRefType>(fifoOut.getElementType());
    int lenOut = elemTypeOut.getNumElements();
    for (size_t i = 0; i < getFifoIns().size(); i++) {
      int len = 0;
      int offset = *getConstantIntValue(getSrcOffsets()[i]);
      if (i == getFifoIns().size() - 1)
        len = lenOut - *getConstantIntValue(getSrcOffsets()[i]);
      else
        len = *getConstantIntValue(getSrcOffsets()[i + 1]) - offset;
      lengths.push_back(len);
    }
  }
  return lengths;
}

std::vector<int> ObjectFifoLinkOp::getDistributeTransferLengths() {
  std::vector<int> lengths;
  if (isDistribute()) {
    auto fifoIn =
        llvm::cast<AIEObjectFifoType>(getInputObjectFifos()[0].getElemType());
    auto elemTypeIn = llvm::cast<MemRefType>(fifoIn.getElementType());
    int lenIn = elemTypeIn.getNumElements();
    for (size_t i = 0; i < getFifoOuts().size(); i++) {
      int offset = *getConstantIntValue(getDstOffsets()[i]);
      int len = 0;
      if (i == getFifoOuts().size() - 1)
        len = lenIn - *getConstantIntValue(getDstOffsets()[i]);
      else
        len = *getConstantIntValue(getDstOffsets()[i + 1]) - offset;
      lengths.push_back(len);
    }
  }
  return lengths;
}

std::optional<int> ObjectFifoLinkOp::getRepeatCount() {
  for (auto fifoOut : getOutputObjectFifos())
    if (fifoOut.getRepeatCount().has_value())
      return {fifoOut.getRepeatCount().value()};
  return {};
}

//===----------------------------------------------------------------------===//
// ObjectFifoRegisterExternalBuffersOp
//===----------------------------------------------------------------------===//

TileOp ObjectFifoRegisterExternalBuffersOp::getTileOp() {
  return cast<TileOp>(getTile().getDefiningOp());
}

ObjectFifoCreateOp ObjectFifoRegisterExternalBuffersOp::getObjectFifo() {
  Operation *parent = getOperation();
  while ((parent = parent->getParentOp())) {
    if (parent->hasTrait<OpTrait::SymbolTable>()) {
      if (auto *st = SymbolTable::lookupSymbolIn(parent, getObjFifoName());
          isa_and_nonnull<ObjectFifoCreateOp>(st))
        return dyn_cast<ObjectFifoCreateOp>(st);
    }
  }
  return {};
}

//===----------------------------------------------------------------------===//
// ObjectFifoAcquireOp
//===----------------------------------------------------------------------===//

LogicalResult ObjectFifoAcquireOp::verify() {
  auto parent = getOperation()->getParentOfType<CoreOp>();
  if (parent == nullptr)
    return emitOpError("must be called from inside a CoreOp");

  auto coreTile = parent.getTile();
  auto objFifo = getObjectFifo();
  if (!objFifo)
    return emitError("cannot retrieve associated object FIFO");
  if (getPort() == ObjectFifoPort::Produce) {
    if (coreTile != objFifo.getProducerTile())
      return parent.emitOpError(
          "producer port of objectFifo accessed by core running "
          "on non-producer tile");
  } else if (getPort() == ObjectFifoPort::Consume) {
    bool found = false;
    for (auto consumerTile : objFifo.getConsumerTiles()) {
      if (coreTile == consumerTile) {
        found = true;
        break;
      }
    }
    if (!found)
      return parent.emitOpError(
          "consumer port of objectFifo accessed by core running "
          "on non-consumer tile");
  }

  auto objFifoElem =
      llvm::cast<AIEObjectFifoType>(getObjectFifo().getElemType())
          .getElementType();
  auto objFifoSubviewElem =
      llvm::cast<AIEObjectFifoSubviewType>(getResult().getType())
          .getElementType();
  // Also accept the consumer element type for asymmetric ObjectFifos
  auto objFifoConsType = llvm::dyn_cast<AIEObjectFifoType>(
      getObjectFifo().getConsumerElemTypeOrDefault());
  if (!objFifoConsType)
    return emitOpError("ObjectFifo consumer element type must be an "
                       "!aie.objectfifo<memref<...>> type");
  auto objFifoConsElem = objFifoConsType.getElementType();
  if (objFifoElem != objFifoSubviewElem &&
      objFifoConsElem != objFifoSubviewElem)
    return emitOpError(
        "ObjectFifo element and ObjectFifoSubview element must match.\n");

  return success();
}

ObjectFifoCreateOp ObjectFifoAcquireOp::getObjectFifo() {
  Operation *parent = getOperation();
  while ((parent = parent->getParentOp())) {
    if (parent->hasTrait<OpTrait::SymbolTable>()) {
      if (auto *st = SymbolTable::lookupSymbolIn(parent, getObjFifoName());
          isa_and_nonnull<ObjectFifoCreateOp>(st))
        return dyn_cast<ObjectFifoCreateOp>(st);
    }
  }
  return {};
}

//===----------------------------------------------------------------------===//
// ObjectFifoReleaseOp
//===----------------------------------------------------------------------===//

LogicalResult ObjectFifoReleaseOp::verify() {
  auto parent = getOperation()->getParentOfType<CoreOp>();
  if (parent == nullptr)
    return emitOpError("must be called from inside a CoreOp");

  auto coreTile = parent.getTile();
  auto objFifo = getObjectFifo();
  if (!objFifo)
    return emitError("cannot retrieve associated object FIFO");
  if (getPort() == ObjectFifoPort::Produce) {
    if (coreTile != objFifo.getProducerTile())
      return parent.emitOpError(
          "producer port of objectFifo accessed by core running "
          "on non-producer tile");
  } else if (getPort() == ObjectFifoPort::Consume) {
    bool found = false;
    for (auto consumerTile : objFifo.getConsumerTiles()) {
      if (coreTile == consumerTile) {
        found = true;
        break;
      }
    }
    if (!found)
      return parent.emitOpError(
          "consumer port of objectFifo accessed by core running "
          "on non-consumer tile");
  }

  return success();
}

ObjectFifoCreateOp ObjectFifoReleaseOp::getObjectFifo() {
  Operation *parent = getOperation();
  while ((parent = parent->getParentOp())) {
    if (parent->hasTrait<OpTrait::SymbolTable>()) {
      if (auto *st = SymbolTable::lookupSymbolIn(parent, getObjFifoName());
          isa_and_nonnull<ObjectFifoCreateOp>(st))
        return dyn_cast<ObjectFifoCreateOp>(st);
    }
  }
  return {};
}

//===----------------------------------------------------------------------===//
// ObjectFifoSubviewAccessOp
//===----------------------------------------------------------------------===//

LogicalResult ObjectFifoSubviewAccessOp::verify() {
  if (auto parent = getOperation()->getParentOfType<CoreOp>();
      parent == nullptr)
    return emitOpError("must be called from inside a CoreOp");

  // The subview operand must be the direct SSA result of an
  // aie.objectfifo.acquire. Flowing it through scf.yield / iter_args / any
  // other region boundary loses the link to the originating acquire that
  // every downstream pass needs (the acquire's size for the bounds check
  // here, and the (fifo, port) identity for lock lowering). Reject here
  // rather than crashing on a null defining op later.
  auto acqOp = getSubview().getDefiningOp<ObjectFifoAcquireOp>();
  if (!acqOp)
    return emitOpError("subview operand must be the direct result of an "
                       "aie.objectfifo.acquire; flowing the subview through "
                       "scf.yield / iter_args / region results is not "
                       "supported");
  if (getIndex() >= acqOp.acqNumber())
    return emitOpError("accessed farther than number of acquired elements "
                       "(index out of bounds).");

  return success();
}

//===----------------------------------------------------------------------===//
// CascadeFlowOp
//===----------------------------------------------------------------------===//

LogicalResult CascadeFlowOp::verify() {
  TileLike src = getSourceTileLike();
  TileLike dst = getDestTileLike();

  if (!src || !dst)
    return emitOpError("source and dest must be tile-like operations");

  if (src.isShimTile() || dst.isShimTile())
    return emitOpError("shimTile row has no cascade stream interface");
  if (src.isMemTile() || dst.isMemTile())
    return emitOpError("memTile row has no cascade stream interface");

  std::optional<int> srcCol = src.tryGetCol();
  std::optional<int> srcRow = src.tryGetRow();
  std::optional<int> dstCol = dst.tryGetCol();
  std::optional<int> dstRow = dst.tryGetRow();

  if (srcCol && srcRow && dstCol && dstRow) {
    const auto &t = getTargetModel(*this);
    if (!t.isSouth(*srcCol, *srcRow, *dstCol, *dstRow) &&
        !t.isWest(*srcCol, *srcRow, *dstCol, *dstRow) &&
        !t.isNorth(*srcCol, *srcRow, *dstCol, *dstRow) &&
        !t.isEast(*srcCol, *srcRow, *dstCol, *dstRow)) {
      return emitOpError("tiles must be adjacent");
    }
  }

  return success();
}

TileLike CascadeFlowOp::getSourceTileLike() {
  return dyn_cast<TileLike>(getSourceTile().getDefiningOp());
}

TileLike CascadeFlowOp::getDestTileLike() {
  return dyn_cast<TileLike>(getDestTile().getDefiningOp());
}

TileOp CascadeFlowOp::getSourceTileOp() {
  if (auto tileOp = dyn_cast_or_null<TileOp>(getSourceTile().getDefiningOp()))
    return tileOp;
  llvm::report_fatal_error("Calling getSourceTileOp requires TileOp.");
}

TileOp CascadeFlowOp::getDestTileOp() {
  if (auto tileOp = dyn_cast_or_null<TileOp>(getDestTile().getDefiningOp()))
    return tileOp;
  llvm::report_fatal_error("Calling getDestTileOp requires TileOp.");
}

//===----------------------------------------------------------------------===//
// ConfigureCascadeOp
//===----------------------------------------------------------------------===//

LogicalResult ConfigureCascadeOp::verify() {
  if (!isa<TileOp>(getTile().getDefiningOp()))
    return emitOpError("requires a placed tile (aie.tile), not a logical tile");

  const auto &t = getTargetModel(*this);
  TileOp tile = cast<TileOp>(getTile().getDefiningOp());
  CascadeDir inputDir = getInputDir();
  CascadeDir outputDir = getOutputDir();

  if (tile.isShimTile())
    return emitOpError("shimTile row has no cascade stream interface");
  if (tile.isMemTile())
    return emitOpError("memTile row has no cascade stream interface");

  if (isa<AIE2TargetModel>(t)) {
    if (inputDir == CascadeDir::South || inputDir == CascadeDir::East) {
      return emitOpError("input direction of cascade must be North or West on ")
             << stringifyAIEArch(t.getTargetArch());
    }
    if (outputDir == CascadeDir::North || outputDir == CascadeDir::West) {
      return emitOpError(
                 "output direction of cascade must be South or East on ")
             << stringifyAIEArch(t.getTargetArch());
    }
  } else {
    return emitOpError("cascade not supported in ")
           << stringifyAIEArch(t.getTargetArch());
  }
  return success();
}

//===----------------------------------------------------------------------===//
// PutCascadeOp
//===----------------------------------------------------------------------===//

LogicalResult PutCascadeOp::verify() {
  const auto &targetModel = getTargetModel(*this);
  Type type = getCascadeValue().getType();
  DataLayout dataLayout = DataLayout::closest(*this);
  auto bits = dataLayout.getTypeSizeInBits(type);
  auto archbits = targetModel.getAccumulatorCascadeSize();
  if (bits != archbits)
    return emitOpError("type must match architecture cascade width (")
           << archbits << " bits in "
           << stringifyAIEArch(targetModel.getTargetArch()) << ")";
  return success();
}

//===----------------------------------------------------------------------===//
// GetCascadeOp
//===----------------------------------------------------------------------===//

LogicalResult GetCascadeOp::verify() {
  const auto &targetModel = getTargetModel(*this);
  Type type = getCascadeValue().getType();
  DataLayout dataLayout = DataLayout::closest(*this);
  auto bits = dataLayout.getTypeSizeInBits(type);
  if (isa<AIE1TargetModel>(targetModel)) {
    if (bits != 384)
      return emitOpError("must be a 384-bit type");
  } else if (isa<AIE2TargetModel>(targetModel)) {
    if (bits != 512)
      return emitOpError("must be a 512-bit type");
  } else
    return emitOpError("cascade not supported in ")
           << stringifyAIEArch(targetModel.getTargetArch());
  return success();
}

//===----------------------------------------------------------------------===//
// DeviceOp
//===----------------------------------------------------------------------===//

LogicalResult DeviceOp::verify() {
  // A compute tile has exactly one core in hardware, so at most one aie.core
  // may resolve to any given (col, row). Cores whose tile is a logical_tile
  // with unspecified coordinates are skipped here — they cannot collide until
  // --aie-place-tiles assigns them a position, at which point this same check
  // runs again on the resulting aie.tile coordinates.
  DenseMap<TileID, CoreOp> coreAtTile;
  WalkResult result = walk([&](CoreOp core) {
    auto tile =
        llvm::dyn_cast_or_null<TileLike>(core.getTile().getDefiningOp());
    if (!tile)
      return WalkResult::advance();
    std::optional<int> col = tile.tryGetCol();
    std::optional<int> row = tile.tryGetRow();
    if (!col || !row)
      return WalkResult::advance();
    TileID id{*col, *row};
    auto [it, inserted] = coreAtTile.try_emplace(id, core);
    if (!inserted) {
      InFlightDiagnostic diag =
          core.emitOpError()
          << "tile (" << *col << ", " << *row
          << ") already has a core; each compute tile can host only one core";
      diag.attachNote(it->second.getLoc()) << "the other core is here";
      return WalkResult::interrupt();
    }
    return WalkResult::advance();
  });
  if (result.wasInterrupted())
    return failure();
  return success();
}

const AIETargetModel &DeviceOp::getTargetModel() {
  return xilinx::AIE::getTargetModel(getDevice());
}

xilinx::AIE::DeviceOp DeviceOp::getForSymbolInModule(mlir::ModuleOp module,
                                                     llvm::StringRef symbol) {
  DeviceOp deviceOp;
  if (!symbol.size()) {
    // If no device name is given, assume 'main'
    symbol = "main";
  }
  Operation *maybeDeviceOp = mlir::SymbolTable::lookupSymbolIn(module, symbol);
  if (!maybeDeviceOp) {
    return nullptr;
  }
  deviceOp = llvm::dyn_cast<DeviceOp>(maybeDeviceOp);
  return deviceOp;
}

xilinx::AIE::DeviceOp
DeviceOp::getForSymbolInModuleOrError(mlir::ModuleOp module,
                                      llvm::StringRef symbol) {
  DeviceOp deviceOp = getForSymbolInModule(module, symbol);
  if (!deviceOp) {
    if (!symbol.empty()) {
      module.emitError("No such device: ") << symbol;
    } else {
      module.emitError("No 'main' device in module");
    }
  }
  return deviceOp;
}

//===----------------------------------------------------------------------===//
// TileElement
//===----------------------------------------------------------------------===//

TileOp TileElement::getTileOp() {
  auto element = cast<TileElement>(this->getOperation());
  Operation *definingOp = element.getTile().getDefiningOp();
  if (auto tileOp = dyn_cast_or_null<TileOp>(definingOp))
    return tileOp;
  llvm::report_fatal_error("Calling getTileOp requires TileOp.");
}

//===----------------------------------------------------------------------===//
// LogicalTileOp
//===----------------------------------------------------------------------===//

LogicalResult LogicalTileOp::verify() {
  const auto &targetModel = getTargetModel(*this);
  int columns = targetModel.columns();
  int rows = targetModel.rows();

  // Only verify col/row bounds if they are specified
  if (auto col = getCol()) {
    if (*col >= columns)
      return emitOpError("column index (")
             << *col
             << ") must be less than the number of columns in the device ("
             << columns << ")";
  }
  if (auto row = getRow()) {
    if (*row >= rows)
      return emitOpError("row index (")
             << *row << ") must be less than the number of rows in the device ("
             << rows << ")";
  }

  // Check that the specified tile type exists on the target device
  AIETileType tileType = getTileType();
  bool tileTypeExists = false;
  for (int col = 0; col < columns && !tileTypeExists; col++) {
    for (int row = 0; row < rows && !tileTypeExists; row++) {
      if (targetModel.getTileType(col, row) == tileType)
        tileTypeExists = true;
    }
  }
  if (!tileTypeExists) {
    return emitOpError("tile type '")
           << stringifyAIETileType(tileType)
           << "' does not exist on the target device";
  }

  // Check logical tile type matches coordinates on device
  // Only validate when both col and row are specified
  if (auto col = tryGetCol()) {
    if (auto row = tryGetRow()) {
      if (targetModel.getTileType(*col, *row) != tileType) {
        return emitOpError("declared logical tile type does not match "
                           "the tile type at coordinates (")
               << *col << ", " << *row << ")";
      }
    }
  }

  if (isShimNOCorPLTile() && getAllocationScheme())
    return emitOpError("Shim tiles cannot have an allocation scheme");

  return success();
}

TileID LogicalTileOp::getCanonicalTileID() {
  const auto &targetModel = getTargetModel(*this);

  // If col and row are both specified, use them directly
  if (getCol().has_value() && getRow().has_value()) {
    return {getCol().value(), getRow().value()};
  }

  // Otherwise, find a representative tile of the given type
  AIETileType tileType = getTileType();
  for (int col = 0; col < targetModel.columns(); col++) {
    for (int row = 0; row < targetModel.rows(); row++) {
      if (targetModel.getTileType(col, row) == tileType) {
        return {col, row};
      }
    }
  }
  llvm_unreachable("No tile of matching tile type found in AIE device");
}

size_t LogicalTileOp::getNumSourceConnections(WireBundle bundle) {
  const auto &targetModel = getTargetModel(*this);
  TileID tile = getCanonicalTileID();

  if (bundle == WireBundle::Core || bundle == WireBundle::DMA) {
    // Note dest is correct here, since direction is reversed.
    if (isShimNOCorPLTile())
      return targetModel.getNumDestShimMuxConnections(tile.col, tile.row,
                                                      bundle);
    return targetModel.getNumDestSwitchboxConnections(tile.col, tile.row,
                                                      bundle);
  }
  return 0;
}

size_t LogicalTileOp::getNumDestConnections(WireBundle bundle) {
  const auto &targetModel = getTargetModel(*this);
  TileID tile = getCanonicalTileID();

  if (bundle == WireBundle::Core || bundle == WireBundle::DMA) {
    // Note source is correct here, since direction is reversed.
    if (isShimNOCorPLTile())
      return targetModel.getNumDestShimMuxConnections(tile.col, tile.row,
                                                      bundle);
    return targetModel.getNumSourceSwitchboxConnections(tile.col, tile.row,
                                                        bundle);
  }
  return 0;
}

std::optional<int> LogicalTileOp::tryGetCol() {
  if (auto col = getCol())
    return col;
  return std::nullopt;
}

std::optional<int> LogicalTileOp::tryGetRow() {
  if (auto row = getRow())
    return row;
  return std::nullopt;
}

//===----------------------------------------------------------------------===//
// Custom Printer and Parser for LogicalTileOp
//===----------------------------------------------------------------------===//

ParseResult LogicalTileOp::parse(OpAsmParser &parser, OperationState &result) {
  AIETileType tileType;
  if (parser.parseLess())
    return failure();

  StringRef tileTypeStr;
  if (parser.parseKeyword(&tileTypeStr))
    return failure();

  auto tileTypeOpt = symbolizeAIETileType(tileTypeStr);
  if (!tileTypeOpt)
    return parser.emitError(parser.getCurrentLocation(),
                            "unknown logical tile type: ")
           << tileTypeStr;
  tileType = *tileTypeOpt;

  if (parser.parseGreater())
    return failure();

  if (parser.parseLParen())
    return failure();

  std::optional<int32_t> col;
  if (succeeded(parser.parseOptionalQuestion())) {
    // col is unspecified
  } else {
    int32_t colVal;
    if (parser.parseInteger(colVal))
      return failure();
    col = colVal;
  }

  if (parser.parseComma())
    return failure();

  std::optional<int32_t> row;
  if (succeeded(parser.parseOptionalQuestion())) {
    // row is unspecified
  } else {
    int32_t rowVal;
    if (parser.parseInteger(rowVal))
      return failure();
    row = rowVal;
  }

  if (parser.parseRParen())
    return failure();

  // Parse optional attributes
  if (parser.parseOptionalAttrDict(result.attributes))
    return failure();

  // Add the parsed attributes to the result
  result.getOrAddProperties<LogicalTileOp::Properties>().tile_type =
      AIETileTypeAttr::get(parser.getContext(), tileType);
  if (col)
    result.getOrAddProperties<LogicalTileOp::Properties>().col =
        parser.getBuilder().getI32IntegerAttr(*col);
  if (row)
    result.getOrAddProperties<LogicalTileOp::Properties>().row =
        parser.getBuilder().getI32IntegerAttr(*row);

  // Add result type (index)
  result.addTypes(parser.getBuilder().getIndexType());

  return success();
}

void LogicalTileOp::print(OpAsmPrinter &printer) {
  printer << "<" << stringifyAIETileType(getTileType()) << ">";

  printer << "(";
  if (auto col = getCol())
    printer << *col;
  else
    printer << "?";
  printer << ", ";
  if (auto row = getRow())
    printer << *row;
  else
    printer << "?";
  printer << ")";

  SmallVector<StringRef, 3> elidedAttrs = {"tile_type", "col", "row"};
  printer.printOptionalAttrDict((*this)->getAttrs(), elidedAttrs);
}

//===----------------------------------------------------------------------===//
// TileOp
//===----------------------------------------------------------------------===//

LogicalResult TileOp::verify() {
  const auto &targetModel = getTargetModel(*this);
  int columns = targetModel.columns();
  int rows = targetModel.rows();
  if (colIndex() >= columns)
    return emitOpError("column index (")
           << colIndex()
           << ") must be less than the number of columns in the device ("
           << columns << ")";
  if (rowIndex() >= rows)
    return emitOpError("row index (")
           << rowIndex()
           << ") must be less than the number of rows in the device (" << rows
           << ")";

  auto users = getResult().getUsers();
  bool found = false;
  for (auto *user : users) {
    if (llvm::isa<SwitchboxOp>(*user)) {
      if (found)
        return emitOpError("can only have one switchbox");
      found = true;
    }
  }

  if (isShimNOCorPLTile() && getAllocationScheme())
    return emitOpError("Shim tiles cannot have an allocation scheme");

  return success();
}

size_t TileOp::getNumSourceConnections(WireBundle bundle) {
  const auto &targetModel = getTargetModel(*this);
  if (bundle == WireBundle::Core || bundle == WireBundle::DMA)
  // Note dest is correct here, since direction is reversed.
  {
    // Note dest is correct here, since direction is reversed.
    if (isShimNOCorPLTile())
      return targetModel.getNumDestShimMuxConnections(getCol(), getRow(),
                                                      bundle);
    return targetModel.getNumDestSwitchboxConnections(getCol(), getRow(),
                                                      bundle);
  }
  return 0;
}

size_t TileOp::getNumDestConnections(WireBundle bundle) {
  const auto &targetModel = getTargetModel(*this);
  if (bundle == WireBundle::Core || bundle == WireBundle::DMA)
  // Note source is correct here, since direction is reversed.
  {
    // Note source is correct here, since direction is reversed.
    if (isShimNOCorPLTile())
      return targetModel.getNumDestShimMuxConnections(getCol(), getRow(),
                                                      bundle);
    return targetModel.getNumSourceSwitchboxConnections(getCol(), getRow(),
                                                        bundle);
  }
  return 0;
}

std::optional<int> TileOp::tryGetCol() { return getCol(); }
std::optional<int> TileOp::tryGetRow() { return getRow(); }

AIETileType TileOp::getTileType() {
  const auto &targetModel = getTargetModel(*this);
  return targetModel.getTileType(getCol(), getRow());
}

bool isLegalTileConnection(TileOp tile, const AIETargetModel &targetModel,
                           MasterSetOp masterOp, PacketRulesOp slaveOp) {
  auto srcBundle = slaveOp.sourcePort().bundle;
  auto srcChan = slaveOp.sourcePort().channel;
  auto dstBundle = masterOp.destPort().bundle;
  auto dstChan = masterOp.destPort().channel;
  return targetModel.isLegalTileConnection(
      tile.colIndex(), tile.rowIndex(), srcBundle, srcChan, dstBundle, dstChan);
}

bool isLegalTileConnection(TileOp tile, const AIETargetModel &targetModel,
                           ConnectOp connectOp) {
  auto srcBundle = connectOp.getSourceBundle();
  auto srcChan = connectOp.getSourceChannel();
  auto dstBundle = connectOp.getDestBundle();
  auto dstChan = connectOp.getDestChannel();
  return targetModel.isLegalTileConnection(
      tile.colIndex(), tile.rowIndex(), srcBundle, srcChan, dstBundle, dstChan);
}

TileOp TileOp::getOrCreate(mlir::OpBuilder builder, DeviceOp device, int col,
                           int row) {
  TileOp tile = nullptr;
  // Find matching predefined tile at device top level, ...
  for (auto t : device.getOps<AIE::TileOp>()) {
    if (t.getRow() == row && t.getCol() == col) {
      tile = t;
      break;
    }
  }
  // ... or if undefined, create a new tile op
  if (!tile) {
    OpBuilder::InsertionGuard guard(builder);
    mlir::Block &device_start_block = *device.getBodyRegion().begin();
    builder.setInsertionPointToStart(&device_start_block);
    tile = TileOp::create(builder, device.getLoc(), builder.getIndexType(), col,
                          row);
  }
  return tile;
}

//===----------------------------------------------------------------------===//
// ShimMuxOp
//===----------------------------------------------------------------------===//

LogicalResult ShimMuxOp::verify() {
  // ShimMux requires a placed tile (TileOp), not a logical tile
  if (!isa<TileOp>(getTile().getDefiningOp()))
    return emitOpError("requires a placed tile (aie.tile), not a logical tile");

  Region &body = getConnections();
  DenseSet<Port> destset;
  if (body.empty())
    return emitOpError("should have non-empty body");

  for (auto &ops : body.front()) {
    if (auto connectOp = dyn_cast<ConnectOp>(ops)) {
      Port dest = {connectOp.getDestBundle(), connectOp.destIndex()};
      if (destset.count(dest))
        return connectOp.emitOpError("targets same destination ")
               << stringifyWireBundle(dest.bundle) << ": " << dest.channel
               << " as another connect operation";
      destset.insert(dest);
    } else if (isa<EndOp>(ops)) {
      // continue;
    } else {
      return ops.emitOpError("cannot be contained in a Switchbox op");
    }
  }
  return success();
}

size_t ShimMuxOp::getNumSourceConnections(WireBundle bundle) {
  auto tile = getTileOp();
  const auto &targetModel = getTargetModel(*this);
  return targetModel.getNumSourceShimMuxConnections(tile.getCol(),
                                                    tile.getRow(), bundle);
}

size_t ShimMuxOp::getNumDestConnections(WireBundle bundle) {
  auto tile = getTileOp();
  const auto &targetModel = getTargetModel(*this);
  return targetModel.getNumDestShimMuxConnections(tile.getCol(), tile.getRow(),
                                                  bundle);
}

//===----------------------------------------------------------------------===//
// ShimDMAOp
//===----------------------------------------------------------------------===//

LogicalResult ShimDMAOp::verify() {
  if (HasSomeTerminator<DMAStartOp, NextBDOp, EndOp>::verifyTrait(*this)
          .failed())
    return failure();
  return success();
}

TileOp ShimDMAOp::getTileOp() {
  return cast<TileElement>(this->getOperation()).getTileOp();
}

LogicalResult PacketRulesOp::verify() {
  if (Region &body = getRules(); body.empty())
    return emitOpError("should have non-empty body");
  return success();
}

LogicalResult PacketFlowOp::verify() {
  Region &body = getPorts();
  if (body.empty())
    return emitOpError("should have non-empty body");

  int numSources = 0, numDests = 0;
  for (auto &ops : body.front()) {
    if (!isa<PacketSourceOp, PacketDestOp, EndOp>(ops))
      return ops.emitOpError("cannot be contained in a PacketFlow op");
    if (isa<PacketSourceOp>(ops))
      ++numSources;
    if (isa<PacketDestOp>(ops))
      ++numDests;
  }

  if (numSources < 1)
    return emitOpError("must have at least one aie.packet_source");
  if (numDests < 1)
    return emitOpError("must have at least one aie.packet_dest");

  return success();
}

//===----------------------------------------------------------------------===//
// CoreOp
//===----------------------------------------------------------------------===//

LogicalResult CoreOp::verify() {
  if (getBody().empty())
    return emitOpError("should have non-empty body");
  if (getElfFile()) {
    // If an ELF file is specified, no MLIR body is allowed (to remove
    // ambiguity); the ELF file will fully dictate what runs on the
    // core and any MLIR would be ignored.
    if (!isEmpty()) {
      return emitOpError(
          "When `elf_file` attribute is specified, core body must be empty "
          "(consist of exactly one `aie.end` op).");
    }
  }
  if (getLinkWith() && getLinkFiles())
    return emitOpError(
        "cannot specify both 'link_with' (deprecated) and 'link_files' "
        "on the same core; run aie-assign-core-link-files to migrate");
  return success();
}

bool CoreOp::isEmpty() {
  Region &body = getBody();
  // Return iff. core body contains exactly one block with exactly one AIE.EndOp
  return (body.hasOneBlock() && body.front().getOperations().size() == 1 &&
          llvm::isa<AIE::EndOp>(body.front().front()));
}

TileOp CoreOp::getTileOp() {
  return cast<TileElement>(this->getOperation()).getTileOp();
}

//===----------------------------------------------------------------------===//
// BufferOp
//===----------------------------------------------------------------------===//

int64_t BufferOp::getAllocationSize() {
  auto type = llvm::cast<MemRefType>(getType());
  DataLayout dataLayout = DataLayout::closest(getOperation());
  return type.getNumElements() * dataLayout.getTypeSize(type.getElementType());
}

LogicalResult BufferOp::verify() {
  if (UsesAreAccessible::verifyTrait(*this).failed())
    return failure();
  return success();
}

// FIXME: make address assignment for buffers explicit and move this function to
// an interface
int32_t xilinx::AIE::getBufferBaseAddress(Operation *bufOp) {
  if (auto buf = dyn_cast<BufferOp>(bufOp)) {
    assert(buf.getAddress().has_value() && "buffer must have address assigned");
    return buf.getAddress().value();
  }
  if (isa_and_nonnull<ExternalBufferOp>(bufOp))
    llvm::report_fatal_error(
        "External buffer addresses are assigned at runtime.");
  llvm::report_fatal_error("unknown buffer type");
}

void xilinx::AIE::collectTiles(DeviceOp &device,
                               DenseMap<TileID, Operation *> &tiles) {
  for (auto tile : device.getOps<TileOp>()) {
    tiles[tile.getTileID()] = tile;
  }
}

void xilinx::AIE::collectBuffers(
    DeviceOp &device,
    DenseMap<Operation *, SmallVector<BufferOp, 4>> &buffers) {
  for (BufferOp buffer : device.getOps<BufferOp>()) {
    Operation *tileOp = buffer.getTile().getDefiningOp();
    buffers[tileOp].push_back(buffer);
  }
}

static void printBufferInitialValue(OpAsmPrinter &p, BufferOp op, Type type,
                                    Attribute initialValue) {
  if (op.getInitialValue()) {
    p << "= ";
    p.printAttributeWithoutType(initialValue);
  }
}

static ParseResult parseBufferInitialValue(OpAsmParser &parser, Type &type,
                                           Attribute &initialValue) {
  auto memrefType = llvm::cast<MemRefType>(type);
  if (!memrefType.hasStaticShape())
    return parser.emitError(parser.getNameLoc())
           << "type should be static shaped memref, but got " << type;

  if (parser.parseOptionalEqual())
    return success();

  Type tensorType = mlir::memref::getTensorTypeFromMemRefType(memrefType);
  if (parser.parseAttribute(initialValue, tensorType))
    return failure();
  if (!llvm::isa<ElementsAttr>(initialValue))
    return parser.emitError(parser.getNameLoc())
           << "initial value should be an elements attribute";
  return success();
}

TileOp BufferOp::getTileOp() {
  return cast<TileElement>(this->getOperation()).getTileOp();
}

//===----------------------------------------------------------------------===//
// MemOp
//===----------------------------------------------------------------------===//

LogicalResult MemOp::verify() {
  Region &body = getBody();
  if (HasSomeTerminator<DMAStartOp, NextBDOp, EndOp>::verifyTrait(*this)
          .failed())
    return failure();

  for (auto &bodyOp : body.getOps()) {
    if (auto allocOp = dyn_cast<memref::AllocOp>(bodyOp))
      if (!allocOp->getAttr("id"))
        return allocOp.emitOpError()
               << "allocOp in MemOp region should have an id attribute";
  }
  return success();
}

TileOp MemOp::getTileOp() {
  return cast<TileElement>(this->getOperation()).getTileOp();
}

//===----------------------------------------------------------------------===//
// MemTileDMAOp
//===----------------------------------------------------------------------===//

LogicalResult MemTileDMAOp::verify() {
  assert(getOperation()->getNumRegions() == 1 &&
         "MemTileDMAOp has zero region!");

  if (HasSomeTerminator<DMAStartOp, NextBDOp, EndOp>::verifyTrait(*this)
          .failed())
    return failure();

  for (auto &bodyOp : getBody().getOps()) {
    if (auto allocOp = dyn_cast<memref::AllocOp>(bodyOp)) {
      if (!allocOp->getAttr("id"))
        return allocOp.emitOpError()
               << "allocOp in MemTileDMAOp region should have an id attribute";
    }
    if (auto startOp = dyn_cast<DMAStartOp>(bodyOp)) {
      if (startOp.getChannelIndex() > 3) {
        // Channels 4 and 5 in a memtile are restricted to only access local
        // buffers and locks.

        // TODO: Move this code to the dialect
        // Set of blocks found to be reachable within a given region.
        llvm::SmallSet<Block *, 16> reachable;
        SmallVector<Block *, 16> worklist;
        Block *firstBD = startOp.getSuccessor(0);
        reachable.insert(firstBD);
        worklist.push_back(firstBD);
        while (!worklist.empty()) {
          Block *block = worklist.pop_back_val();
          if (block->empty())
            continue;
          auto successors = block->getTerminator()->getSuccessors();
          for (auto *i : successors) {
            if (!reachable.contains(i)) {
              reachable.insert(i);
              worklist.push_back(i);
            }
          }
        }
        for (Block *b : reachable) {
          for (DMABDOp bd : b->getOps<DMABDOp>()) {
            if (auto bufferOp = bd.getBufferOp();
                bufferOp.getTile() != getTile()) {
              InFlightDiagnostic err =
                  bd.emitOpError()
                  << "is reachable from DMA channel "
                  << startOp.getChannelIndex()
                  << " and attempts to access a non-local buffer\n";
              err.attachNote(startOp->getLoc()) << "channel";
              err.attachNote(bufferOp->getLoc()) << "buffer";
              return err;
            }
          }
          for (auto useLock : b->getOps<UseLockOp>()) {
            if (auto lockOp = useLock.getLockOp();
                lockOp.getTile() != getTile()) {
              InFlightDiagnostic err =
                  useLock.emitOpError()
                  << "is reachable from DMA channel "
                  << startOp.getChannelIndex()
                  << " and attempts to access a non-local lock\n";
              err.attachNote(startOp->getLoc()) << "channel";
              err.attachNote(lockOp->getLoc()) << "lock";
              return err;
            }
          }
        }
      }
    }
  }

  return success();
}

TileOp MemTileDMAOp::getTileOp() {
  return cast<TileElement>(this->getOperation()).getTileOp();
}

//===----------------------------------------------------------------------===//
// DMAOp
//===----------------------------------------------------------------------===//

LogicalResult DMAOp::verify() {
  auto *parentOp = getOperation()->getParentOp();
  if (parentOp->getRegion(0).getBlocks().size() > 1)
    return emitOpError("DMAOp can only appear in single block region");
  if (!parentOp->getRegion(0).getOps<DMAStartOp>().empty())
    return emitOpError("DMAOp is not compatible with DMAStart ops");
  auto bdRegions = getBds();
  for (auto &bdRegion : bdRegions) {
    if (!bdRegion.hasOneBlock())
      return emitOpError("DMAOp regions must have only one block");
    auto bds = llvm::to_vector_of<DMABDOp>(bdRegion.front().getOps<DMABDOp>());
    if (bds.size() != 1)
      return emitOpError("DMAOp regions/blocks must have exactly one DMABDOp");
    auto useLocks =
        llvm::to_vector_of<UseLockOp>(bdRegion.front().getOps<UseLockOp>());
    if (useLocks.size() != 2)
      return emitOpError(
          "DMAOp regions/blocks must have exactly two UseLock ops");
  }
  return success();
}

//===----------------------------------------------------------------------===//
// DMABDOp
//===----------------------------------------------------------------------===//

BufferOp DMABDOp::getBufferOp() {
  return cast<BufferOp>(getBuffer().getDefiningOp());
}

// let assemblyFormat = [{
//   `(` $buffer `:` type($buffer) (`,` $offset^)? (`,` $len^)? (`,`
//   $dimensions^)? (`,` $pad_dimensions^)? (`,` `pad_value` `=` $pad_value^)?
//   `)` attr-dict
// }];
ParseResult DMABDOp::parse(OpAsmParser &parser, OperationState &result) {
  OpAsmParser::UnresolvedOperand bufferRawOperand{};
  ::llvm::ArrayRef<OpAsmParser::UnresolvedOperand> bufferOperands(
      &bufferRawOperand, 1);
  ::llvm::SMLoc bufferOperandsLoc;
  (void)bufferOperandsLoc;
  Type bufferRawType{};
  ::llvm::ArrayRef<Type> bufferTypes(&bufferRawType, 1);
  IntegerAttr offsetAttr;
  IntegerAttr lenAttr;
  ::xilinx::AIE::BDDimLayoutArrayAttr dimensionsAttr;
  ::xilinx::AIE::BDPadLayoutArrayAttr pad_dimensionsAttr;
  IntegerAttr pad_valueAttr;
  if (parser.parseLParen())
    return failure();

  bufferOperandsLoc = parser.getCurrentLocation();
  if (parser.parseOperand(bufferRawOperand))
    return failure();
  if (parser.parseColon())
    return failure();
  if (parser.parseCustomTypeWithFallback(bufferRawType))
    return failure();

  // offset
  if (succeeded(parser.parseOptionalComma())) {
    if (parser.parseCustomAttributeWithFallback(
            offsetAttr, parser.getBuilder().getIntegerType(32))) {
      return failure();
    }
    if (!offsetAttr)
      offsetAttr = parser.getBuilder().getIntegerAttr(
          parser.getBuilder().getIntegerType(32), 0);
    result.getOrAddProperties<DMABDOp::Properties>().offset = offsetAttr;
  }

  // len
  if (succeeded(parser.parseOptionalComma())) {
    if (parser.parseCustomAttributeWithFallback(
            lenAttr, parser.getBuilder().getIntegerType(32))) {
      return failure();
    }
    if (lenAttr)
      result.getOrAddProperties<DMABDOp::Properties>().len = lenAttr;
  }

  // dimensions
  if (succeeded(parser.parseOptionalComma())) {
    if (parser.parseCustomAttributeWithFallback(dimensionsAttr, Type{})) {
      return failure();
    }
    if (dimensionsAttr)
      result.getOrAddProperties<DMABDOp::Properties>().dimensions =
          dimensionsAttr;
  }

  // pad_dimensions
  if (succeeded(parser.parseOptionalComma())) {
    if (parser.parseCustomAttributeWithFallback(pad_dimensionsAttr, Type{})) {
      return failure();
    }
    if (pad_dimensionsAttr)
      result.getOrAddProperties<DMABDOp::Properties>().pad_dimensions =
          pad_dimensionsAttr;
  }

  // pad_value
  if (succeeded(parser.parseOptionalComma())) {
    if (parser.parseKeyword("pad_value"))
      return failure();
    if (parser.parseEqual())
      return failure();

    if (parser.parseCustomAttributeWithFallback(
            pad_valueAttr, parser.getBuilder().getIntegerType(32))) {
      return failure();
    }
    if (pad_valueAttr)
      result.getOrAddProperties<DMABDOp::Properties>().pad_value =
          pad_valueAttr;
  }
  if (parser.parseRParen())
    return failure();

  auto loc = parser.getCurrentLocation();
  if (parser.parseOptionalAttrDict(result.attributes))
    return failure();
  if (failed(verifyInherentAttrs(result.name, result.attributes, [&]() {
        return parser.emitError(loc)
               << "'" << result.name.getStringRef() << "' op ";
      })))
    return failure();

  if (parser.resolveOperands(bufferOperands, bufferTypes, bufferOperandsLoc,
                             result.operands))
    return failure();

  return success();
}

void DMABDOp::print(::mlir::OpAsmPrinter &printer) {
  printer << "(";
  printer << getBuffer();
  printer << ' ' << ":";
  printer << ' ';
  {
    auto type = getBuffer().getType();
    if (auto validType = ::llvm::dyn_cast<::mlir::MemRefType>(type))
      printer.printStrippedAttrOrType(validType);
    else
      printer << type;
  }
  if (getLenAttr() ||
      getOffsetAttr() !=
          ::mlir::OpBuilder((*this)->getContext())
              .getIntegerAttr(
                  ::mlir::OpBuilder((*this)->getContext()).getIntegerType(32),
                  0)) {
    printer << ",";
    printer << ' ';
    printer.printAttributeWithoutType(getOffsetAttr());
  }
  if (getLenAttr()) {
    printer << ",";
    printer << ' ';
    printer.printAttributeWithoutType(getLenAttr());
  }
  if (getDimensionsAttr()) {
    printer << ",";
    printer << ' ';
    printer.printStrippedAttrOrType(getDimensionsAttr());
  }
  if (getPadDimensionsAttr()) {
    printer << ",";
    printer << ' ';
    printer.printStrippedAttrOrType(getPadDimensionsAttr());
  }
  if ((getPadValueAttr() &&
       getPadValueAttr() !=
           ::mlir::OpBuilder((*this)->getContext())
               .getIntegerAttr(
                   ::mlir::OpBuilder((*this)->getContext()).getIntegerType(32),
                   0))) {
    printer << ",";
    printer << ' ' << "pad_value";
    printer << ' ' << "=";
    printer << ' ';
    printer.printAttributeWithoutType(getPadValueAttr());
  }
  printer << ")";
  ::llvm::SmallVector<::llvm::StringRef, 2> elidedAttrs;
  elidedAttrs.push_back("offset");
  elidedAttrs.push_back("len");
  elidedAttrs.push_back("dimensions");
  elidedAttrs.push_back("pad_dimensions");
  elidedAttrs.push_back("pad_value");
  printer.printOptionalAttrDict((*this)->getAttrs(), elidedAttrs);
}

// A BDDimLayoutAttr array (outermost-first) describes a contiguous row-major
// scan when the innermost stride is 1 and each outer stride equals the product
// of all inner sizes.  Used by both DMABDOp verification and canonicalization.
bool xilinx::AIE::isContiguousBDTransfer(llvm::ArrayRef<BDDimLayoutAttr> dims) {
  if (dims.empty())
    return true; // no ND layout = trivially contiguous
  // Innermost (last) stride must be 1.
  if (dims.back().getStride() != 1)
    return false;
  // Each outer stride must equal the product of all inner sizes.
  // Use uint64_t throughout to match the unsigned getSize()/getStride() types.
  uint64_t product = 1;
  for (int i = static_cast<int>(dims.size()) - 1; i >= 1; --i) {
    product *= dims[i].getSize();
    // A size-1 dimension's stride is irrelevant (it is never stepped).
    if (dims[i - 1].getSize() > 1 && dims[i - 1].getStride() != product)
      return false;
  }
  return true;
}

LogicalResult DMABDOp::verify() {
  // Skip verification of the BDOp outside of mem operations.
  // BDOps may appear elsewhere and subsequent lowerings will place them in the
  // correct mem ops.
  Operation *p = (*this)->getParentOp();
  if (!llvm::isa<MemOp, MemTileDMAOp, ShimDMAOp, DMAOp>(*p)) {
    return success();
  }

  // Check if buffer is an unranked memref (e.g., from function argument)
  bool isUnrankedMemRef = llvm::isa<UnrankedMemRefType>(getBuffer().getType());

  // For unranked memrefs, we can't verify as strictly since we don't know
  // the buffer's defining op or its full type at compile time
  if (!isUnrankedMemRef) {
    if (!isa<BufferOp, ExternalBufferOp>(getBuffer().getDefiningOp()))
      return emitOpError(
          "BDs only support BufferOp or ExternalBufferOp operands.");
  }

  if (getLenInBytes() % 4)
    return emitOpError("transfer length must be multiple of 4 (i.e., represent "
                       "4 byte aligned address)");

  TileElement parentTileElement = getParentTileElement(getOperation());
  TileLike parentTile = parentTileElement.getTileLike();
  if (!parentTile)
    return emitOpError("parent tile must implement TileLike interface");

  if (!isUnrankedMemRef && getOperation()->getParentOfType<MemOp>() &&
      getBufferOp().getTile() != parentTileElement.getTile())
    return emitOpError(
        "Core tile DMAs can only access a buffer in the same tile.");

  const AIETargetModel &targetModel = getTargetModel(getOperation());

  uint32_t maxBds = targetModel.getNumBDs(parentTile.getTileType());
  if (std::optional<int32_t> bdId = getBdId();
      bdId.has_value() && static_cast<uint32_t>(*bdId) >= maxBds)
    return emitOpError("bdId attribute exceeds max: ") << maxBds - 1;
  if (std::optional<int32_t> nextBdId = getNextBdId();
      nextBdId.has_value() && static_cast<uint32_t>(*nextBdId) >= maxBds)
    return emitOpError("nextBdId attribute exceeds max: ") << maxBds - 1;
  if (auto dims = getDimensions(); dims.has_value()) {
    size_t maxNDims = 3;
    if (getOperation()->getParentOfType<MemTileDMAOp>())
      maxNDims = 4;
    if (dims->size() > maxNDims)
      return emitOpError() << "Cannot give more than "
                           << std::to_string(maxNDims)
                           << " dimensions for step sizes and wraps in this "
                              " tile (got "
                           << std::to_string(dims->size()) << " dimensions).";

    auto buffer = llvm::dyn_cast<MemRefType>(getBuffer().getType());
    if (!buffer)
      return emitOpError() << "dimensions attribute cannot be used with "
                              "unranked memref buffer type.";
    int64_t maxIdx = getDimsMaxIdx(*dims);
    if (buffer.getNumElements() <= maxIdx)
      return emitOpError() << "Specified stride(s) and size(s) result in out "
                              "of bounds access in buffer, for index "
                           << std::to_string(maxIdx) << " in memref of length "
                           << std::to_string(buffer.getNumElements()) << ".";

    // A contiguous row-major access on a shim tile is lowered to linear mode
    // by aie-dma-tasks-to-npu / aie-dma-to-npu, using the wide buffer_length
    // register which is exempt from the 10-bit ND wrap-size limit.
    // Skip the per-dimension size check when the BD is on a shim tile and the
    // access is contiguous, so the natural ND form can be written without
    // triggering a spurious verifier error before lowering.
    //
    // Note: the verifier early-exit above means we only reach this code when
    // the parent op is MemOp, MemTileDMAOp, ShimDMAOp, or DMAOp -- all of
    // which are TileElements, so parentTile is always non-null here.
    bool skipSizeCheck =
        parentTile.isShimTile() && xilinx::AIE::isContiguousBDTransfer(*dims);

    for (BDDimLayoutAttr dim : *dims) {
      if (0 == dim.getStride())
        return emitOpError()
               << "Invalid step size; must be a positive integer.";
      if (dim.getStride() > buffer.getNumElements())
        return emitOpError() << "Step size " << std::to_string(dim.getStride())
                             << " exceeds memref size "
                             << std::to_string(buffer.getNumElements());
      if (!skipSizeCheck && dim.getSize() >= (1UL << 9) + 1)
        return emitOpError() << "Size may not exceed 1023.";
      if (dim.getStride() >= (1UL << 19))
        return emitOpError() << "Stride may not exceed " << (1 << 20);
    }

    // Since streams read 32b words, there's no way to read eg 16b with stride
    // of 2 (ie lower halfs of each 32b). So force it to be 1 (and then in
    // CDODirect/XAIEV2 scale the size by 4/getBufferElementTypeWidthInBytes).
    if (getBufferElementTypeWidthInBytes() < 4 && dims->back().getStride() != 1)
      return emitOpError(
          "For <32b width datatypes, inner-most dim stride must be 1");

    if (getBufferElementTypeWidthInBytes() > 4 && dims->back().getStride() != 1)
      return emitOpError(
          "For >32b width datatypes, inner-most dim stride must be 1");
  }
  if (auto paddims = getPadDimensions(); paddims.has_value()) {
    auto dims = getDimensions();
    if (!dims.has_value())
      return emitOpError() << "Padding requires n-d data layouts expressed as"
                           << " wrap(s) and stride(s).";
    if (!parentTile.isMemTile())
      return emitOpError() << "Padding is only supported by memtile dma bds.";
    if (dims->size() != paddims->size())
      return emitOpError() << "Mismatch number of dimensions between padding(s)"
                           << " and wrap(s) and stride(s).";
    int actuallen = 1;
    for (unsigned i = 0; i < paddims->size(); i++) {
      auto dim = (*dims)[i];
      auto paddim = (*paddims)[i];
      actuallen *= paddim.getConstPadBefore() + paddim.getConstPadAfter() +
                   dim.getSize();
      if (actuallen > getLen())
        return emitOpError() << "Data exceeds len after padding.";
    }
    if ((paddims->back().getConstPadBefore() *
         getBufferElementTypeWidthInBytes()) %
        4)
      return emitOpError() << "Inner-most padding-before count must result in"
                           << " padding in 32-bit words.";
    if ((paddims->back().getConstPadAfter() *
         getBufferElementTypeWidthInBytes()) %
        4)
      return emitOpError() << "Inner-most padding-after count must result in"
                           << " padding in 32-bit words.";
  }
  if (!isUnrankedMemRef &&
      (parentTile.isMemTile() || parentTile.isCoreTile())) {
    if (auto baseAddr = getBufferOp().getAddress(); baseAddr.has_value()) {
      int64_t offsetInBytes = *baseAddr + getOffsetInBytes();
      if (offsetInBytes % 4)
        return emitOpError("bd address must be 4 byte (32b) aligned; got "
                           "base+offset: ")
               << offsetInBytes << " (bytes)";
    }
  }
  if (auto packetInfo = getPacket()) {
    if (packetInfo->getPktType() > 7)
      return emitOpError("Packet type field can only hold 3 bits.");
    if (packetInfo->getPktId() > 31)
      return emitOpError("Packet ID field can only hold 5 bits.");
  }

  if (!getLen() && !getBuffer().getType().hasStaticShape())
    return emitOpError() << "buffer with dynamic shape requires static length.";

  if (getBurstLength() != 0 && !parentTile.isShimNOCTile())
    return emitOpError("Burst length is only supported in Shim NOC tiles that "
                       "are connected to the memory-mapped NOC.");

  return success();
}

//===----------------------------------------------------------------------===//
// DMAStartOp
//===----------------------------------------------------------------------===//

static LogicalResult FoldDMAStartOp(DMAStartOp op, PatternRewriter &rewriter) {

  llvm::SetVector<Block *> reachable;
  SmallVector<Block *, 16> worklist;
  Block *firstBD = op.getSuccessor(0);
  reachable.insert(firstBD);
  worklist.push_back(firstBD);
  while (!worklist.empty()) {
    Block *block = worklist.pop_back_val();
    if (block->empty())
      continue;
    auto successors = block->getTerminator()->getSuccessors();
    for (auto *i : successors) {
      if (!reachable.contains(i)) {
        reachable.insert(i);
        worklist.push_back(i);
      }
    }
  }

  // BD chain ends with an EndOp, indicating non-repeating pattern: BD chain
  // folding not applicable.
  if (isa<EndOp>((reachable.back())->getTerminator()))
    return failure();

  // Check for identical bds.
  auto areEquivalentBDs = [](Block *b1, Block *b2) {
    auto b1OpRange = b1->without_terminator();
    auto b2OpRange = b2->without_terminator();
    if (llvm::range_size(b1OpRange) != llvm::range_size(b2OpRange))
      return false;
    if (!llvm::all_of(llvm::zip_equal(b1OpRange, b2OpRange),
                      [](std::tuple<Operation &, Operation &> pair) {
                        return OperationEquivalence::isEquivalentTo(
                            &std::get<0>(pair), &std::get<1>(pair),
                            OperationEquivalence::IgnoreLocations);
                      }))
      return false;
    return true;
  };

  // Get a vector of unique BDs.
  SmallVector<Block *> uniquePattern;
  auto patternIt = reachable.begin();
  while (patternIt != reachable.end() &&
         llvm::none_of(uniquePattern, [patternIt, areEquivalentBDs](Block *b1) {
           return areEquivalentBDs(*patternIt, b1);
         })) {
    uniquePattern.push_back(*patternIt);
    patternIt++;
  }

  unsigned idx = 0;
  while (patternIt != reachable.end()) {
    // BD repetition found. Check if repeating pattern.
    if (!areEquivalentBDs(*patternIt, uniquePattern[idx]))
      return failure();
    patternIt++;
    idx = (++idx) % uniquePattern.size();
  }

  // Repeating BD chains detected. Erasing repetitions.
  auto lastBDTerm = dyn_cast<NextBDOp>(reachable.back()->getTerminator());
  auto lastUniqueBDTerm =
      dyn_cast<NextBDOp>(uniquePattern.back()->getTerminator());
  lastUniqueBDTerm.setSuccessor(lastBDTerm.getSuccessor());

  return success();
}

// Canonicalization pattern for DMABDOp: on shim tiles, fold a contiguous
// row-major ND access pattern into canonical linear form (no dimensions
// attribute), so that the hardware uses the wide buffer_length register.
namespace {
struct LinearizeContiguousBDTransfer : public mlir::OpRewritePattern<DMABDOp> {
  using OpRewritePattern::OpRewritePattern;

  mlir::LogicalResult
  matchAndRewrite(DMABDOp op, mlir::PatternRewriter &rewriter) const override {
    // Only fire for shim DMA BDs: ExternalBufferOp buffer
    // (dma_configure_task_for stage), parent TileElement is shim
    // (dma_configure_task stage), or inside ShimDMAOp (after full lowering).
    bool bufferIsExternal =
        isa_and_nonnull<ExternalBufferOp>(op.getBuffer().getDefiningOp());
    TileElement parentElem = getParentTileElement(op.getOperation());
    TileLike parentTileLike =
        parentElem ? parentElem.getTileLike() : TileLike{};
    bool parentIsShim = parentTileLike && parentTileLike.isShimTile();
    bool inShimDMA = (bool)op->getParentOfType<ShimDMAOp>();
    if (!bufferIsExternal && !parentIsShim && !inShimDMA)
      return mlir::failure();

    // Only when ND dimensions are present and contiguous.
    auto dims = op.getDimensions();
    if (!dims.has_value() || dims->empty())
      return mlir::failure();
    if (!xilinx::AIE::isContiguousBDTransfer(*dims))
      return mlir::failure();
    // Already linear (single dimension with stride 1)?
    if (dims->size() == 1 && dims->front().getStride() == 1)
      return mlir::failure();

    int64_t product = 1;
    for (BDDimLayoutAttr dim : *dims)
      product *= dim.getSize();

    // len < product: the outermost dim describes a hardware BD iteration
    // (preserved downstream as iteration_size/stride). Don't fold it away.
    if (auto lenVal = op.getLen())
      if (static_cast<int64_t>(*lenVal) != product)
        return mlir::failure();
    int32_t len = static_cast<int32_t>(product);

    // Drop the dimensions attribute in-place; all other attributes (offset,
    // len, packet, burst_length, bd_id, etc.) are preserved automatically.
    rewriter.modifyOpInPlace(op, [&]() {
      op.setLen(len);
      op->removeAttr("dimensions");
    });
    return mlir::success();
  }
};
} // namespace

void DMABDOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                          MLIRContext *context) {
  patterns.add<LinearizeContiguousBDTransfer>(context);
}

void DMAStartOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                             MLIRContext *context) {
  patterns.add(FoldDMAStartOp);
}

//===----------------------------------------------------------------------===//
// SwitchboxOp
//===----------------------------------------------------------------------===//

LogicalResult SwitchboxOp::verify() {
  // Switchbox requires a placed tile (TileOp), not a logical tile
  if (!isa<TileOp>(getTile().getDefiningOp()))
    return emitOpError("requires a placed tile (aie.tile), not a logical tile");

  Region &body = getConnections();
  DenseSet<Port> sourceset;
  DenseSet<Port> destset;
  auto tile = getTileOp();
  const auto &targetModel = getTargetModel(tile);
  if (body.empty())
    return emitOpError("should have non-empty body");
  for (auto &ops : body.front()) {
    // Would be simpler if this could be templatized.
    auto checkBound = [&ops](StringRef dir, WireBundle bundle, int index,
                             int bound) -> LogicalResult {
      if (index >= bound) {
        if (bound > 0)
          return ops.emitOpError("index ")
                 << index << " for " << dir << " bundle "
                 << stringifyWireBundle(bundle) << " must be less than "
                 << bound;
        return ops.emitOpError()
               << dir << " bundle " << stringifyWireBundle(bundle)
               << " not supported; index: " << index << ", bound: " << bound;
      }
      return success();
    };

    if (auto connectOp = dyn_cast<ConnectOp>(ops)) {
      Port source = {connectOp.getSourceBundle(), connectOp.sourceIndex()};
      sourceset.insert(source);

      Port dest = {connectOp.getDestBundle(), connectOp.destIndex()};
      if (destset.count(dest)) {
        return connectOp.emitOpError()
               << "; connecting " << to_string(source) << " to "
               << to_string(dest) << " on "
               << to_string(this->getTileOp().getTileID())
               << " targets same dst as another connect op; existing "
                  "destinations: "
               << llvm::join(llvm::map_range(
                                 destset, [](auto &p) { return to_string(p); }),
                             ", ");
      }
      destset.insert(dest);

      if (connectOp.sourceIndex() < 0)
        return connectOp.emitOpError("source index cannot be less than zero");

      if (checkBound("source", connectOp.getSourceBundle(),
                     connectOp.sourceIndex(),
                     getNumSourceConnections(connectOp.getSourceBundle()))
              .failed())
        return failure();

      if (connectOp.destIndex() < 0)
        return connectOp.emitOpError("dest index cannot be less than zero");

      if (checkBound("dest", connectOp.getDestBundle(), connectOp.destIndex(),
                     getNumDestConnections(connectOp.getDestBundle()))
              .failed())
        return failure();

      // Stream switch connection constraints
      if (!isLegalTileConnection(tile, targetModel, connectOp))
        return connectOp.emitOpError("illegal stream switch connection");

    } else if (auto connectOp = dyn_cast<MasterSetOp>(ops)) {
      Port dest = {connectOp.getDestBundle(), connectOp.destIndex()};
      if (destset.count(dest))
        return connectOp.emitOpError("targets same destination ")
               << stringifyWireBundle(dest.bundle) << ": " << dest.channel
               << " as another connect or masterset operation";
      destset.insert(dest);

      if (connectOp.destIndex() < 0)
        return connectOp.emitOpError("dest index cannot be less than zero");

      if (checkBound("dest", connectOp.getDestBundle(), connectOp.destIndex(),
                     getNumDestConnections(connectOp.getDestBundle()))
              .failed())
        return failure();

      int arbiter = -1;
      for (auto val : connectOp.getAmsels()) {
        auto amsel = dyn_cast<AMSelOp>(val.getDefiningOp());
        if (arbiter != -1 && arbiter != amsel.arbiterIndex())
          return connectOp.emitOpError(
              "a master port can only be tied to one arbiter");
        arbiter = amsel.arbiterIndex();
      }
    } else if (auto connectOp = dyn_cast<PacketRulesOp>(ops)) {
      Port source = {connectOp.getSourceBundle(), connectOp.sourceIndex()};
      if (sourceset.count(source))
        return connectOp.emitOpError("packet switched source ")
               << stringifyWireBundle(source.bundle) << source.channel
               << " cannot match another connect or masterset operation";
      sourceset.insert(source);

    } else if (auto amselOp = dyn_cast<AMSelOp>(ops)) {
      std::vector<MasterSetOp> mstrs;
      std::vector<PacketRulesOp> slvs;
      for (auto *user : amselOp.getResult().getUsers()) {
        if (auto s = dyn_cast<PacketRuleOp>(user)) {
          auto pktRules = dyn_cast<PacketRulesOp>(s->getParentOp());
          slvs.push_back(pktRules);
        } else if (auto m = dyn_cast<MasterSetOp>(user))
          mstrs.push_back(m);
      }
      for (auto m : mstrs) {
        for (auto s : slvs) {
          // Stream switch connection constraints
          if (!isLegalTileConnection(tile, targetModel, m, s)) {
            return amselOp->emitOpError("illegal stream switch connection");
          }
        }
      }
    } else if (isa<EndOp>(ops)) {
      // continue;
    } else {
      return ops.emitOpError("cannot be contained in a Switchbox op");
    }
  }

  return success();
}

template <typename... ParentOpTypes>
struct HasSomeParent {
  static LogicalResult verifyTrait(Operation *op) {
    Operation *operation = op->getParentOp();
    while (operation) {
      if (llvm::isa_and_nonnull<ParentOpTypes...>(operation))
        return success();
      operation = operation->getParentOp();
    }
    return failure();
  }
};

TileOp LockOp::getTileOp() {
  return cast<TileElement>(this->getOperation()).getTileOp();
}

LogicalResult LockOp::verify() {
  if (auto result = UsesAreAccessible::verifyTrait(*this); result.failed())
    return result;

  if (getLockID().has_value()) {
    TileLike tileLike = getTileLike();
    if (!tileLike)
      return emitOpError("tile operand must implement TileLike interface");
    const auto &targetModel = getTargetModel(*this);
    auto tileType = tileLike.getTileType();
    if (int numLocks = targetModel.getNumLocks(tileType);
        getLockID().value() >= numLocks)
      return emitOpError("lock assigned invalid id (maximum is ")
             << numLocks - 1 << ")";
  }

  return success();
}

// Look up for compile-time constant lock values, if any.
// Returns std::nullopt if lock value does not reference an `arith.constant`.
static std::optional<int32_t> getConstantLockValue(UseLockOp op) {
  if (auto constant = op.getValue().getDefiningOp<arith::ConstantOp>())
    if (auto intAttr = llvm::dyn_cast<IntegerAttr>(constant.getValue()))
      return (int32_t)intAttr.getInt();
  return std::nullopt;
}

struct UsesOneLockInDMABlock {
  static LogicalResult verifyTrait(Operation *op) {
    auto *block = op->getBlock();
    int lockID = -1;
    for (auto op : block->getOps<UseLockOp>()) {
      if (auto lock = dyn_cast<LockOp>(op.getLock().getDefiningOp());
          lock.getLockID().has_value()) {
        if (lockID != -1 && lockID != lock.getLockIDValue())
          return failure();
        lockID = lock.getLockIDValue();
      }
    }
    return success();
  }
};

struct AcquireReleaseOneStateInDMABlock {
  static LogicalResult verifyTrait(Operation *op) {
    auto *block = op->getBlock();
    int acqValue = -1, relValue = -1;
    for (auto op : block->getOps<UseLockOp>()) {
      // Non-constant lock values cannot be compared here; the passes that
      // require a constant enforce that separately via getConstantValue().
      auto value = getConstantLockValue(op);
      if (!value)
        continue;
      if (op.acquire() || op.acquireGE()) {
        if (acqValue != -1 && acqValue != *value) {
          return failure();
        }
        acqValue = *value;
      } else if (op.release()) {
        if (relValue != -1 && relValue != *value) {
          return failure();
        }
        relValue = *value;
      }
    }
    return success();
  }
};

struct AccessesLocalLocks {
  static LogicalResult verifyTrait(Operation *op) {
    if (auto memOp = op->getParentOfType<MemOp>()) {
      auto useLock = dyn_cast<UseLockOp>(op);
      if (auto lock = useLock.getLockOp(); lock.getTile() != memOp.getTile())
        return failure();
    }
    return success();
  }
};

LogicalResult UseLockOp::verify() {
  // AIE.useLock cannot be used at the top level
  if (llvm::isa_and_nonnull<DeviceOp, ModuleOp>((*this)->getParentOp()))
    return (*this)->emitOpError("must be used in a core or memory operation.");

  const auto &targetModel = getTargetModel(*this);
  if (targetModel.getTargetArch() == AIEArch::AIE1 && acquireGE())
    return (*this)->emitOpError(
        "AcquireGreaterEqual is not supported in AIE1.");

  // Locks used inside a DMA/BD block are configured via static register writes
  // and therefore require a compile-time constant value.
  if (HasSomeParent<MemOp, MemTileDMAOp, ShimDMAOp>::verifyTrait(*this)
          .succeeded() &&
      !getConstantLockValue(*this))
    return (*this)->emitOpError(
        "lock value in a DMA/BD block must be a compile-time constant "
        "(defined by an arith.constant).");

  // Otherwise, AIE.useLock should be inside MemOp, MemTileDMAOp, or
  // ShimDMAOp,
  if (HasSomeParent<MemOp, MemTileDMAOp, ShimDMAOp>::verifyTrait(*this)
          .succeeded()) {
    if (!(*this)->getBlock())
      return (*this)->emitOpError("is not in a block.");

    if (targetModel.getTargetArch() == AIEArch::AIE1 &&
        UsesOneLockInDMABlock::verifyTrait(*this).failed())
      return (*this)->emitOpError(
          "used in a DMA block that have multiple locks.");

    if (AcquireReleaseOneStateInDMABlock::verifyTrait(*this).failed())
      return (*this)->emitOpError("acquires/releases the lock in a DMA block "
                                  "from/to multiple states.");

    if (HasSomeParent<MemOp>::verifyTrait(*this).succeeded() &&
        AccessesLocalLocks::verifyTrait(*this).failed())
      return (*this)->emitOpError("can only access a lock in the same tile");
    return success();

    // Or it can be in a CoreOp, or some FuncOp called from a CoreOp
  }
  if (HasSomeParent<CoreOp, func::FuncOp>::verifyTrait(*this).succeeded()) {
    return success();
  }
  // Or it can be in a DMAConfigureTaskOp (for runtime DMA configuration)
  // Check by operation name to avoid circular dependency with AIEX dialect
  {
    Operation *operation = (*this)->getParentOp();
    while (operation) {
      if (operation->getName().getStringRef() == "aiex.dma_configure_task")
        return success();
      operation = operation->getParentOp();
    }
  }
  return (*this)->emitOpError()
         << "expects some parent op to be one of "
         << "AIE::device, AIE::core, func::func, AIE::mem, AIE::shimDMA, or "
            "AIEX::dma_configure_task";
}

#include "aie/Dialect/AIE/IR/AIEEnums.cpp.inc"
#include "aie/Dialect/AIE/IR/AIEInterfaces.cpp.inc"

//===----------------------------------------------------------------------===//
// TraceEventAttr
//===----------------------------------------------------------------------===//

// Custom parser for TraceEventAttr value (uses shared helper)
static ParseResult parseTraceEventValue(AsmParser &parser, Attribute &value) {
  return xilinx::AIE::parseTraceEvent(parser, value);
}

// Custom printer for TraceEventAttr value (uses shared helper)
static void printTraceEventValue(AsmPrinter &printer, Attribute value) {
  xilinx::AIE::printTraceEventEnum(printer, value);
}

// Custom parser for TraceEventAttr value (uses shared helper)
static ParseResult parseTraceRegValue(OpAsmParser &parser, Attribute &value) {

  // Try to parse as number
  int64_t intValue;
  OptionalParseResult parseResult = parser.parseOptionalInteger(intValue);
  if (parseResult.has_value() && succeeded(parseResult.value())) {
    MLIRContext *ctx = parser.getContext();
    value = IntegerAttr::get(IntegerType::get(ctx, 32), intValue);
    return success();
  }
  return xilinx::AIE::parseTraceEvent(parser, value);
}

// Custom printer for TraceEventAttr value (uses shared helper)
static void printTraceRegValue(OpAsmPrinter &printer, Operation *op,
                               Attribute value) {
  // if it's an intattr
  if (auto intAttr = llvm::dyn_cast<IntegerAttr>(value)) {
    printer << intAttr.getInt();
    return;
  }
  xilinx::AIE::printTraceEventEnum(printer, value);
}

// Helper to parse a LockBlocking enum keyword.
#define GET_OP_CLASSES
#include "aie/Dialect/AIE/IR/AIEOps.cpp.inc"

FailureOr<int32_t> UseLockOp::getConstantValue() {
  if (auto value = getConstantLockValue(*this))
    return *value;
  return emitOpError("expected the lock value to be a compile-time constant "
                     "(defined by an arith.constant).");
}

TileOp SwitchboxOp::getTileOp() {
  return cast<TileElement>(this->getOperation()).getTileOp();
}

size_t SwitchboxOp::getNumSourceConnections(WireBundle bundle) {
  auto tile = getTileOp();
  const auto &targetModel = getTargetModel(*this);
  return targetModel.getNumSourceSwitchboxConnections(tile.getCol(),
                                                      tile.getRow(), bundle);
}

size_t SwitchboxOp::getNumDestConnections(WireBundle bundle) {
  auto tile = getTileOp();
  const auto &targetModel = getTargetModel(*this);
  return targetModel.getNumDestSwitchboxConnections(tile.getCol(),
                                                    tile.getRow(), bundle);
}

TileOp ShimMuxOp::getTileOp() {
  return cast<TileElement>(this->getOperation()).getTileOp();
}

WireBundle xilinx::AIE::getConnectingBundle(WireBundle dir) {
  switch (dir) {
  case WireBundle::North:
    return WireBundle::South;
  case WireBundle::South:
    return WireBundle::North;
  case WireBundle::East:
    return WireBundle::West;
  case WireBundle::West:
    return WireBundle::East;
  default:
    return dir;
  }
}

//===----------------------------------------------------------------------===//
// BDChainOp
//===----------------------------------------------------------------------===//

ParseResult BDChainOp::parse(OpAsmParser &parser, OperationState &result) {
  SmallVector<OpAsmParser::Argument> entryArgs;

  // Symbol name, e.g. @my_chain
  StringAttr symNameAttr;
  if (parser.parseSymbolName(symNameAttr, SymbolTable::getSymbolAttrName(),
                             result.attributes)) {
    return failure();
  }

  // Entry arguments (placeholders), e.g. (%addr: memref<1xi32>)
  ParseResult argParseResult = parser.parseCommaSeparatedList(
      OpAsmParser::Delimiter::Paren, [&]() -> ParseResult {
        OpAsmParser::Argument argument;
        if (parser.parseArgument(argument, true, true)) {
          return failure();
        }
        entryArgs.push_back(argument);
        return success();
      });
  if (argParseResult) {
    return argParseResult;
  }

  // BD Chain Body
  auto *body = result.addRegion();
  ParseResult bodyParseResult = parser.parseRegion(*body, entryArgs, false);
  if (bodyParseResult) {
    return bodyParseResult;
  }

  return success();
}

void BDChainOp::print(OpAsmPrinter &printer) {
  auto taskName =
      (*this)
          ->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName())
          .getValue();
  printer << ' ';
  printer.printSymbolName(taskName);

  Region &body = getRegion();
  auto argsIter = body.getArguments();
  printer << '(';
  for (auto it = argsIter.begin(); it != argsIter.end(); ++it) {
    if (it != argsIter.begin()) {
      printer << ", ";
    }
    printer.printRegionArgument(*it);
  }
  printer << ')';

  printer << ' ';
  printer.printRegion(body, false, true);
}

//===----------------------------------------------------------------------===//
// ShimDMAAllocationOp
//===----------------------------------------------------------------------===//

LogicalResult ShimDMAAllocationOp::verify() {
  TileLike tileLike = llvm::dyn_cast<TileLike>(getTile().getDefiningOp());
  if (!tileLike) {
    return emitOpError("tile operand must implement TileLike interface");
  }

  if (!tileLike.isShimNOCorPLTile()) {
    // if placed, provide detailed error message
    auto col = tileLike.tryGetCol();
    auto row = tileLike.tryGetRow();
    if (col && row) {
      return emitOpError("tile must be a shim tile, but got tile(")
             << *col << ", " << *row << ")";
    }
    return emitOpError("tile must be a shim tile");
  }

  return success();
}

TileOp ShimDMAAllocationOp::getTileOp() {
  return cast<TileOp>(getTile().getDefiningOp());
}

ShimDMAAllocationOp ShimDMAAllocationOp::getForSymbol(DeviceOp device,
                                                      llvm::StringRef symbol) {
  Operation *maybeOp = device.lookupSymbol(symbol);
  if (maybeOp) {
    if (ShimDMAAllocationOp op = dyn_cast<ShimDMAAllocationOp>(maybeOp)) {
      return op;
    }
  }
  return nullptr;
}

//===----------------------------------------------------------------------===//
// RuntimeSequenceOp
//===----------------------------------------------------------------------===//

ParseResult RuntimeSequenceOp::parse(OpAsmParser &parser,
                                     OperationState &result) {

  // Name of this runtime sequence
  StringAttr nameAttr;
  (void)parser.parseOptionalSymbolName(
      nameAttr, mlir::SymbolTable::getSymbolAttrName(), result.attributes);

  SmallVector<OpAsmParser::Argument> entryArgs;

  // Entry arguments,  e.g. (%addr: memref<1xi32>)
  ParseResult argParseResult = parser.parseCommaSeparatedList(
      OpAsmParser::Delimiter::Paren, [&]() -> ParseResult {
        OpAsmParser::Argument argument;
        if (parser.parseArgument(argument, true, true)) {
          return failure();
        }
        entryArgs.push_back(argument);
        return success();
      });
  if (argParseResult) {
    return argParseResult;
  }

  // Optional `attributes { ... }` clause (before the body so it does not
  // conflict with the `{` that opens the region).
  if (parser.parseOptionalAttrDictWithKeyword(result.attributes))
    return failure();

  // Body
  auto *body = result.addRegion();
  ParseResult bodyParseResult = parser.parseRegion(*body, entryArgs, false);
  if (bodyParseResult) {
    return bodyParseResult;
  }

  return success();
}

void RuntimeSequenceOp::print(OpAsmPrinter &printer) {
  Region &body = getRegion();

  auto nameAttr = (*this)->getAttrOfType<StringAttr>(
      mlir::SymbolTable::getSymbolAttrName());
  if (nameAttr &&
      nameAttr != ::mlir::OpBuilder((*this)->getContext())
                      .getStringAttr(getDefaultRuntimeSequenceName())) {
    printer << ' ';
    printer.printSymbolName(nameAttr);
  }

  printer << '(';
  for (unsigned i = 0, n = body.getNumArguments(); i < n; i++) {
    if (i > 0) {
      printer << ", ";
    }
    printer.printRegionArgument(body.getArgument(i));
  }
  printer << ')';

  printer.printOptionalAttrDictWithKeyword(
      (*this)->getAttrs(),
      /*elidedAttrs=*/{mlir::SymbolTable::getSymbolAttrName()});

  printer << ' ';
  printer.printRegion(body, false, true);
}

LogicalResult RuntimeSequenceOp::verify() {
  DeviceOp device = (*this)->getParentOfType<DeviceOp>();
  if (!device) {
    // this check is redudnant with the HasParent trait, but can't hurt
    (*this)->emitOpError() << "must be inside AIE device operation.";
    return failure();
  }
  return success();
}

RuntimeSequenceOp
RuntimeSequenceOp::getForSymbolInDevice(DeviceOp deviceOp,
                                        llvm::StringRef symbol) {
  RuntimeSequenceOp runtimeSequenceOp;
  if (!symbol.size()) {
    runtimeSequenceOp = *deviceOp.getOps<RuntimeSequenceOp>().begin();
  } else {
    Operation *maybeRuntimeSequenceOp =
        mlir::SymbolTable::lookupSymbolIn(deviceOp, symbol);
    if (!maybeRuntimeSequenceOp) {
      return nullptr;
    }
    runtimeSequenceOp =
        llvm::dyn_cast<RuntimeSequenceOp>(maybeRuntimeSequenceOp);
  }
  return runtimeSequenceOp;
}

RuntimeSequenceOp
RuntimeSequenceOp::getForSymbolInDeviceOrError(DeviceOp deviceOp,
                                               llvm::StringRef symbol) {
  RuntimeSequenceOp runtimeSequenceOp = getForSymbolInDevice(deviceOp, symbol);
  if (!runtimeSequenceOp) {
    if (!symbol.empty()) {
      deviceOp.emitError("No such runtime sequence: ") << symbol;
    } else {
      deviceOp.emitError("No runtime sequence in device");
    }
  }
  return runtimeSequenceOp;
}

LogicalResult RuntimeSequenceOp::verifyBeforeMaterialization() {
  // Check that all symbol references within the runtime sequence
  // are either to ShimDMAAllocationOp, DeviceOp or another RuntimeSequenceOp;
  // these are the only symbols that can be lowered with the NPU passes
  auto result = (*this)->walk([&](Operation *op) {
    for (NamedAttribute namedAttr : op->getAttrs()) {
      Attribute attr = namedAttr.getValue();
      auto walkResult = attr.walk([&](SymbolRefAttr symbolRef) {
        Operation *symbolDefOp =
            SymbolTable::lookupNearestSymbolFrom(*this, symbolRef);
        if (symbolDefOp) {
          if (!llvm::isa<ShimDMAAllocationOp>(symbolDefOp) &&
              !llvm::isa<DeviceOp>(symbolDefOp) &&
              !llvm::isa<RuntimeSequenceOp>(symbolDefOp) &&
              !llvm::isa<BufferOp>(symbolDefOp) &&
              !llvm::isa<memref::GlobalOp>(symbolDefOp)) {
            op->emitOpError()
                << "references symbol '"
                << symbolRef.getRootReference().getValue()
                << "' which must be either a ShimDMAAllocationOp, DeviceOp, "
                   "RuntimeSequenceOp, BufferOp, or GlobalOp, "
                   "but got: "
                << symbolDefOp->getName().getStringRef();
            return WalkResult::interrupt();
          }
          if (BufferOp bufferOp = llvm::dyn_cast<BufferOp>(symbolDefOp)) {
            if (!bufferOp.getAddress()) {
              op->emitOpError()
                  << "Unallocated buffer; fixed addresses are required before "
                     "runtime sequence materialization.";
              return WalkResult::interrupt();
            }
          }
        }
        return WalkResult::advance();
      });
      if (walkResult.wasInterrupted()) {
        return WalkResult::interrupt();
      }
    }
    return WalkResult::advance();
  });

  if (result.wasInterrupted()) {
    return failure();
  }

  return success();
}

// Include implementations for custom attributes
#define GET_ATTRDEF_CLASSES
#include "aie/Dialect/AIE/IR/AIEAttrs.cpp.inc"