AIEXDialect.cpp

//===- AIEXDialect.cpp ------------------------------------------*- C++ -*-===//
//
// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
// (c) Copyright 2019 Xilinx Inc.
//
//===----------------------------------------------------------------------===//

#include "aie/Dialect/AIEX/IR/AIEXDialect.h"

#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/DialectImplementation.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Interfaces/DataLayoutInterfaces.h"
#include "mlir/Interfaces/FoldInterfaces.h"
#include "mlir/Transforms/InliningUtils.h"

#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/TypeSize.h"

#include <cstdint>
#include <numeric>

using namespace mlir;
using namespace xilinx;

#include "aie/Dialect/AIEX/IR/AIEXDialect.cpp.inc"

#include "aie/Dialect/AIEX/IR/AIEXEnums.cpp.inc"

#define GET_TYPEDEF_CLASSES
#include "aie/Dialect/AIEX/IR/AIEXTypes.cpp.inc"

namespace xilinx::AIEX {

// FIXME: use Tablegen'd dialect class
void AIEXDialect::initialize() {
  addOperations<
#define GET_OP_LIST
#include "aie/Dialect/AIEX/IR/AIEX.cpp.inc"
      >();
  addTypes<
#define GET_TYPEDEF_LIST
#include "aie/Dialect/AIEX/IR/AIEXTypes.cpp.inc"
      >();
}

} // namespace xilinx::AIEX

#define GET_OP_CLASSES
#include "aie/Dialect/AIEX/IR/AIEX.cpp.inc"

/* Return the correct values to write to the hardware registers to configure
  strides and wraps given the input user-facing strides and wraps.

  In the IR, we express strides in units of element data type, but the hardware
  requires it in units of address granularity. Address granularity currently is
  4 bytes for all hardware.


  User-facing strides/wraps relate to hardware as follows:

   - By default, stride 0 and size 1 is assumed if unspecified.
   - If only N strides/wraps are defined, those define the lowest N dimensions.

  inputStride[3]        == iteration_stride / elemSizeFac + 1
  inputWrap[3]          == iteration_size + 1
    Highest-dimension stride/wrap is iteration count / iteration stride.
  inputStride[2]        == d2_stride / elemSizeFac + 1
                           Note: d2_size is not specified in hardware as it is
                           implicit from the total buffer transfer length
  inputStride[1]        == d1_stride / elemSizeFac + 1
  inputSize[1]          == d1_size
  inputStride[0]        == d0_stride / elemSizeFac + 1
  inputSize[0]          == d0_size / elemSizeFac

  where elemSizeFac == bufferElementSize / addressGranularity
  where bufferElementSize == size in bytes of elements in buffer,
                             e.g. 4 for int32
  where addressGranularity == transfer granularity in hardware, which is
                              4 bytes for all current hardware

  Note: strides are expressed offset by one from user input strides, because the
  hardware does not support a 0 stride (repeat).
  */
void AIEX::getHardwareStridesWraps(const AIE::AIETargetModel &targetModel,
                                   mlir::Operation *op,
                                   mlir::BaseMemRefType referencedBufType,
                                   llvm::SmallVector<int64_t, 4> inputSizes,
                                   llvm::SmallVector<int64_t, 4> inputStrides,
                                   llvm::SmallVector<int64_t, 4> &sizes,
                                   llvm::SmallVector<int64_t, 4> &strides) {
  assert(inputSizes.size() == inputStrides.size());
  assert(sizes.size() == 4);
  assert(strides.size() == 4);

  DataLayout dataLayout = DataLayout::closest(op);
  auto elemWidth =
      dataLayout.getTypeSizeInBits(referencedBufType.getElementType());
  auto addressGranularity = targetModel.getAddressGenGranularity();

  // Output strides and sizes are default-initialized to 0
  std::fill(sizes.begin(), sizes.end(), 0);
  std::fill(strides.begin(), strides.end(), 0);

  if (inputSizes[0] == 0) {
    // Illegal input, this won't transfer anything at all.
    // Leave it to the verification functions to complain to the user.
    return;
  }

  // d0_size, d0_stride
  sizes[0] = inputSizes[0] * elemWidth / addressGranularity;
  if (inputStrides[0] * elemWidth < addressGranularity ||
      (elemWidth > addressGranularity)) {
    // First check:
    // While the hardware cannot transfer less than addressGranularity bits at
    // a time, the user may expresses a contiguous transfer of multiple
    // elements with a stride smaller than addressGranularity. We can thus set
    // the stride to 1 (encoded in hardware as 0) here to allow such transfers.
    // The verification function should ensure that
    //    inputStrides[0] * elemWidth < addressGranularity
    //    iff. inputSize[0] * elemWidth > addressGranularity.
    // Second check:
    // If the element width is larger than addressGranularity, we need to make
    // sure that all bytes are properly copied and therefore the stride must be
    // set to 1 (encoded in hardware as 0).
    // The verification function should ensure that
    //     inputStrides[0] * elemWidth % addressGranularity == 0
    //     && inputStrides[0] == 1 if elemWidth > addressGranularity
    // This makes it impossible to have a stride greater than 1 for
    // elemWidths bigger than addressGranularity, even if they are a multiple of
    // it. Such operations should make use of an additional dimension instead.
    strides[0] = 0;
  } else {
    strides[0] = inputStrides[0] * elemWidth / addressGranularity - 1;
  }

  // d1_size, d1_stride
  sizes[1] = inputSizes[1];
  if (inputSizes[1] > 1) {
    // Stride only matters if we have more than one iteration.
    strides[1] = inputStrides[1] * elemWidth / addressGranularity - 1;
  }

  // d2_size, d2_stride
  sizes[2] = inputSizes[2];
  if (inputSizes[2] > 1) {
    // Stride only matters if we have more than one iteration.
    strides[2] = inputStrides[2] * elemWidth / addressGranularity - 1;
  }

  // iteration_size, iteration_stride
  if (inputSizes[3] > 1) {
    // Stride only matters if we have more than one iteration.
    sizes[3] = inputSizes[3] - 1;
    // Note that the iteration_stride must be positive, just like the other
    // dimensions. However, one can encode a zero-stride "repeat" of the same
    // transfer by setting a positive repeat_count on the pushToQueue instr,
    // and setting the size here to 1. This causes the BD to "wrap" at every
    // single iteration, effectively never adding the specified stride, in turn
    // equalling a repeat without stride.
    if (inputStrides[3] > 0) {
      strides[3] = inputStrides[3] * elemWidth / addressGranularity - 1;
    }
  }
}

mlir::LogicalResult
AIEX::verifyStridesWraps(mlir::Operation *forOp,
                         mlir::BaseMemRefType referencedBufType, int tileCol,
                         int tileRow, llvm::SmallVector<int64_t, 4> inputSizes,
                         llvm::SmallVector<int64_t, 4> inputStrides,
                         llvm::SmallVector<int64_t, 4> hardwareSizes,
                         llvm::SmallVector<int64_t, 4> hardwareStrides,
                         bool skipTransformationChecks) {
  const auto &targetModel = AIE::getTargetModel(forOp);
  auto addressGranularity = targetModel.getAddressGenGranularity();
  DataLayout dataLayout = DataLayout::closest(forOp);
  auto elemWidth =
      dataLayout.getTypeSizeInBits(referencedBufType.getElementType());

  uint32_t wrap_bits = 0;
  uint32_t step_bits = 0;
  uint32_t iter_bits = 6;
  if (targetModel.isShimNOCTile(tileCol, tileRow)) {
    step_bits = 20; // XAIEMLGBL_NOC_MODULE_DMA_BD0_3_D0_STEPSIZE_WIDTH
    wrap_bits = 10; // XAIEMLGBL_NOC_MODULE_DMA_BD0_3_D0_WRAP_WIDTH
  } else if (targetModel.isMemTile(tileCol, tileRow)) {
    step_bits = 17; // XAIEMLGBL_MEM_TILE_MODULE_DMA_BD0_2_D0_STEPSIZE_WIDTH
    wrap_bits = 10; // XAIEMLGBL_MEM_TILE_MODULE_DMA_BD0_2_D0_WRAP_WIDTH
  } else if (targetModel.isCoreTile(tileCol, tileRow)) {
    step_bits = 13; // XAIEMLGBL_MEMORY_MODULE_DMA_BD0_2_D0_STEPSIZE_WIDTH
    wrap_bits = 8;  // XAIEMLGBL_MEMORY_MODULE_DMA_BD0_3_D0_WRAP_WIDTH
  } else {
    return forOp->emitOpError(
        "Unsupported tile type at (" + std::to_string(tileCol) + ", " +
        std::to_string(tileRow) + ") Must be ShimNOC, Mem or Core.");
  }

  for (int i = 0; i < 4; i++) {
    if (inputSizes[i] <= 0) {
      return forOp->emitOpError("Size ") << i << " must be a positive integer.";
    }
  }

  if (inputSizes[0] * elemWidth % addressGranularity != 0) {
    std::stringstream msg;
    msg << "Transfer sizes must be multiples of " << (addressGranularity / 8)
        << " bytes. " << inputSizes[0] << " elements at " << (elemWidth / 8)
        << " bytes each equal " << (inputSizes[0] * elemWidth / 8)
        << " bytes, which is not divisible by " << (addressGranularity / 8)
        << ". ";
    return forOp->emitOpError(msg.str());
  }

  for (int i = 0; i < 3; i++) {
    if (inputSizes[i] > 1 && inputStrides[i] < 1) {
      // If inputSize[i] == 1, anything is allowable in the stride, since that
      // stride will never be applied. For any larger size, we must verify that
      // the stride is positive.
      return forOp->emitOpError("Stride ")
             << i << " must be a positive integer.";
    }
  }
  // A value of zero is allowable for the fourth-dimension stride
  // (this indicates an interation stride for the repeat of 0)
  if (inputSizes[3] > 1 && inputStrides[3] < 0) {
    return forOp->emitOpError("Stride 3 must be a non-negative integer.");
  }

  for (int i = 0; i < 4; i++) {
    // strides[0] == 1 is ok iff the transfer size is a multiple of
    // addressGranularity, which is checked below
    if (i == 0 && inputStrides[i] == 1)
      continue;
    if (inputStrides[i] * elemWidth % addressGranularity != 0) {
      std::stringstream msg;
      msg << "Stride " << i << " is " << inputStrides[i] << " elements * "
          << (elemWidth / 8) << " bytes = " << (inputStrides[i] * elemWidth / 8)
          << " bytes, which is not divisible by " << (addressGranularity / 8)
          << ". ";
      return forOp->emitOpError(msg.str());
    }
  }

  if (!skipTransformationChecks && hardwareSizes[0] > (1 << wrap_bits) - 1)
    return forOp->emitOpError(
        "Size 0 exceeds the [0:" + std::to_string((1 << wrap_bits) - 1) +
        "] range.");
  if (!skipTransformationChecks && hardwareSizes[1] > (1 << wrap_bits) - 1)
    return forOp->emitOpError(
        "Size 1 exceeds the [0:" + std::to_string((1 << wrap_bits) - 1) +
        "] range.");
  if (hardwareSizes[3] > (1 << iter_bits))
    return forOp->emitOpError(
        "Size 3 exceeds the [1:" + std::to_string(1 << iter_bits) + "] range.");
  if (hardwareStrides[0] > (1 << step_bits))
    return forOp->emitOpError("Stride 0 exceeds the [1:" +
                              std::to_string(1 << step_bits) + "] range.");
  if (hardwareStrides[1] > (1 << step_bits))
    return forOp->emitOpError("Stride 1 exceeds the [1:" +
                              std::to_string(1 << step_bits) + "] range.");
  if (hardwareStrides[2] > (1 << step_bits))
    return forOp->emitOpError("Stride 2 exceeds the [1:" +
                              std::to_string(1 << step_bits) + "] range.");
  // strides[3] exceeding the range is ok iff the sizes[3] is one, which is
  // checked below
  if (hardwareStrides[3] > (1 << step_bits) && hardwareSizes[3] > 0)
    return forOp->emitOpError("Stride 3 exceeds the [1:" +
                              std::to_string(1 << step_bits) + "] range.");

  return success();
}

//===----------------------------------------------------------------------===//
// UseTokenOp
//===----------------------------------------------------------------------===//

LogicalResult AIEX::UseTokenOp::verify() {
  auto *parentOp = (*this)->getParentOp();
  if (isa<func::FuncOp>(parentOp) || isa<AIE::CoreOp>(parentOp) ||
      isa<AIE::MemOp>(parentOp) || isa<AIE::ShimDMAOp>(parentOp))
    return success();
  return failure();
}

//===----------------------------------------------------------------------===//
// MulticastOp
//===----------------------------------------------------------------------===//

LogicalResult AIEX::MulticastOp::verify() {
  Region &body = getPorts();
  assert(getOperation()->getNumRegions());
  assert(!body.empty());
  for (auto &ops : body.front())
    if (!isa<MultiDestOp, AIE::EndOp>(ops))
      return ops.emitOpError("cannot be contained in a Multicast op");

  return success();
}

//===----------------------------------------------------------------------===//
// BroadcastPacketOp
//===----------------------------------------------------------------------===//

LogicalResult AIEX::BroadcastPacketOp::verify() {
  Region &body = getPorts();
  assert(getOperation()->getNumRegions());
  assert(!body.empty());
  for (auto &ops : body.front())
    if (!isa<BPIDOp, AIE::EndOp>(ops))
      return ops.emitOpError("cannot be contained in a BroadcastPacket op");

  return success();
}

//===----------------------------------------------------------------------===//
// NpuDmaMemcpyNdOp
//===----------------------------------------------------------------------===//

/* Calculates the offset value to be written to the
 */
int64_t AIEX::NpuDmaMemcpyNdOp::getOffsetInBytes() {
  llvm::SmallVector<int64_t, 4> offsets =
      llvm::map_to_vector(llvm::reverse(getMixedOffsets()), [](OpFoldResult s) {
        return getConstantIntValue(s).value();
      });
  llvm::SmallVector<int64_t, 4> strides =
      llvm::map_to_vector(llvm::reverse(getMixedStrides()), [](OpFoldResult s) {
        return getConstantIntValue(s).value();
      });
  size_t offset = 0;
  size_t R = offsets.size();
  size_t el_bit_width = getElementTypeBitwidth();
  assert(el_bit_width % 8 == 0 &&
         "Expected Memref element bitwidth to be multiple of 8.");
  size_t S = el_bit_width / 8;
  for (size_t i = 0; i < R; i++)
    offset += offsets[i] * strides[i] * S;
  return offset;
}

// Returns true when sizes/strides describe a plain contiguous transfer with
// no data layout transformation (d1/d2 sizes == 1, d0 stride == 1).
// d3 (repeat) is intentionally excluded.
bool AIEX::isLinearTransfer(llvm::ArrayRef<int64_t> sizes,
                            llvm::ArrayRef<int64_t> strides) {
  return sizes[1] == 1 && sizes[2] == 1 && strides[0] == 1 && strides[1] == 0 &&
         strides[2] == 0;
}

// Returns true when sizes/strides (innermost-first) describe a contiguous
// row-major scan: innermost stride == 1 and each outer stride equals the
// product of all inner sizes.  The repeat dimension (index 3) is excluded.
// Size-1 dimensions are allowed to carry any stride value because that stride
// is never applied during the transfer (the loop runs only once).
// This is the vector-form counterpart of AIE::isContiguousBDTransfer.
bool AIEX::isContiguousTransfer(llvm::ArrayRef<int64_t> sizes,
                                llvm::ArrayRef<int64_t> strides) {
  if (strides[0] != 1)
    return false;
  if (sizes[1] > 1 && strides[1] != sizes[0])
    return false;
  if (sizes[2] > 1 && strides[2] != sizes[0] * sizes[1])
    return false;
  return true;
}

// dma_memcpy_nd transfers of the form [*, 1, 1, len][*, 0, 0, 1] do not
// specify any data layout transformation, but simply express a contiguous
// transfer of `len`. The 4th dimension is excluded because a repeat count
// is still compatible with a linear transfer.
bool AIEX::NpuDmaMemcpyNdOp::isLinearTransferWithoutTransformation() {
  llvm::SmallVector<int64_t, 4> inputSizes =
      llvm::map_to_vector(llvm::reverse(getMixedSizes()), [](OpFoldResult s) {
        return getConstantIntValue(s).value();
      });
  llvm::SmallVector<int64_t, 4> inputStrides =
      llvm::map_to_vector(llvm::reverse(getMixedStrides()), [](OpFoldResult s) {
        return getConstantIntValue(s).value();
      });
  return isLinearTransfer(inputSizes, inputStrides);
}

// Canonicalization pattern: rewrite a contiguous row-major access pattern to
// the canonical linear form [s3, 1, 1, N][st3, 0, 0, 1].
//
// Using outermost-first index notation (matching the IR syntax), a 4D access
// [s3, s2, s1, s0][st3, st2, st1, st0] is a contiguous linear scan when:
//   st0 == 1
//   s1 == 1  ||  st1 == s0          (stride irrelevant when size is 1)
//   s2 == 1  ||  st2 == s0 * s1
// yielding a total of N = s0 * s1 * s2 contiguous elements.  The repeat
// dimension s3 / stride st3 is unchanged by the fold.
//
// Note: this pattern applies only to NpuDmaMemcpyNdOp.  The analogous
// dimensionsToStream / dimensionsFromStreamPerConsumer attributes on
// ObjectFifoCreateOp are not canonicalized here; they are lowered separately
// by the ObjectFifo stateful transform pass.
//
// This fold is always semantically valid and never introduces new hardware
// limit violations: in the resulting linear form, isLinearTransferWithout-
// Transformation() returns true, so verifyStridesWraps() skips the 10-bit
// d0 wrap-size check.  The hardware uses a wider buffer_length register in
// linear mode (32-bit on shim tiles, 17-bit on mem tiles, 14-bit on core
// tiles), so N can be much larger than the 10-bit ND wrap limit.
namespace {
struct LinearizeContiguousTransfer
    : public mlir::OpRewritePattern<AIEX::NpuDmaMemcpyNdOp> {
  using OpRewritePattern::OpRewritePattern;

  mlir::LogicalResult
  matchAndRewrite(AIEX::NpuDmaMemcpyNdOp op,
                  mlir::PatternRewriter &rewriter) const override {
    // Only constant sizes/strides/offsets can be analysed statically.
    if (!llvm::all_of(op.getMixedSizes(), [](mlir::OpFoldResult s) {
          return mlir::getConstantIntValue(s).has_value();
        }))
      return mlir::failure();
    if (!llvm::all_of(op.getMixedStrides(), [](mlir::OpFoldResult s) {
          return mlir::getConstantIntValue(s).has_value();
        }))
      return mlir::failure();
    if (!llvm::all_of(op.getMixedOffsets(), [](mlir::OpFoldResult s) {
          return mlir::getConstantIntValue(s).has_value();
        }))
      return mlir::failure();

    // Skip ops that are already in canonical linear form.
    if (op.isLinearTransferWithoutTransformation())
      return mlir::failure();

    // getMixedSizes/Strides/Offsets return outermost-first; reverse to
    // innermost-first so index 0 = d0 (innermost) and index 3 = repeat.
    llvm::SmallVector<int64_t, 4> sizes = llvm::map_to_vector(
        llvm::reverse(op.getMixedSizes()), [](mlir::OpFoldResult s) {
          return mlir::getConstantIntValue(s).value();
        });
    llvm::SmallVector<int64_t, 4> strides = llvm::map_to_vector(
        llvm::reverse(op.getMixedStrides()), [](mlir::OpFoldResult s) {
          return mlir::getConstantIntValue(s).value();
        });
    llvm::SmallVector<int64_t, 4> offsets = llvm::map_to_vector(
        llvm::reverse(op.getMixedOffsets()), [](mlir::OpFoldResult s) {
          return mlir::getConstantIntValue(s).value();
        });

    // Require a contiguous row-major scan.
    if (!AIEX::isContiguousTransfer(sizes, strides))
      return mlir::failure();

    // Fold d0/d1/d2 into one linear count; keep the repeat dimension intact.
    // Build directly in outermost-first order for the replacement op.
    int64_t N = sizes[0] * sizes[1] * sizes[2];
    llvm::SmallVector<int64_t, 4> newSizesOuter = {sizes[3], 1, 1, N};
    llvm::SmallVector<int64_t, 4> newStridesOuter = {strides[3], 0, 0, 1};

    // getOffsetInBytes() computes: sum(offsets[i] * strides[i] * elemSize).
    // After folding, the intermediate strides become 0, so any non-zero offset
    // in those dimensions would silently contribute 0 bytes.  Preserve the
    // correct start address by collapsing the innermost three offset/stride
    // pairs into a single linear element index at d0.
    int64_t linearOffset = offsets[0] * strides[0] + offsets[1] * strides[1] +
                           offsets[2] * strides[2];
    llvm::SmallVector<int64_t, 4> newOffsetsOuter = {offsets[3], 0, 0,
                                                     linearOffset};

    rewriter.replaceOpWithNewOp<AIEX::NpuDmaMemcpyNdOp>(
        op, op.getMemref(),
        /*offsets=*/mlir::ValueRange{},
        /*sizes=*/mlir::ValueRange{},
        /*strides=*/mlir::ValueRange{},
        mlir::DenseI64ArrayAttr::get(op.getContext(), newOffsetsOuter),
        mlir::DenseI64ArrayAttr::get(op.getContext(), newSizesOuter),
        mlir::DenseI64ArrayAttr::get(op.getContext(), newStridesOuter),
        op.getPacketAttr(), op.getMetadata(), op.getIdAttr(),
        op.getIssueTokenAttr(), op.getD0ZeroBeforeAttr(),
        op.getD1ZeroBeforeAttr(), op.getD2ZeroBeforeAttr(),
        op.getD0ZeroAfterAttr(), op.getD1ZeroAfterAttr(),
        op.getD2ZeroAfterAttr(), op.getBurstLengthAttr(),
        op.getOffsetParameterAttr());
    return mlir::success();
  }
};
} // namespace

void AIEX::NpuDmaMemcpyNdOp::getCanonicalizationPatterns(
    mlir::RewritePatternSet &patterns, mlir::MLIRContext *context) {
  patterns.add<LinearizeContiguousTransfer>(context);
}

// Helper method to check if a requested burst length is supported by the target
// model. Returns an error message if the burst length is not supported or an
// empty option otherwise.
static std::optional<std::string>
checkBurstLength(const xilinx::AIE::AIETargetModel &targetModel,
                 uint32_t requestedBurstLength) {
  if (requestedBurstLength != 0) {
    auto bel = targetModel.getShimBurstEncodingsAndLengths();
    auto pair = std::find_if(bel.begin(), bel.end(),
                             [=](const std::pair<uint32_t, uint32_t> &p) {
                               return p.second == requestedBurstLength;
                             });

    if (pair == bel.end()) {
      std::string errorMessage =
          "Requested burst length is not supported by the target. "
          "Supported burst lengths:";

      errorMessage =
          std::accumulate(bel.begin(), bel.end(), errorMessage,
                          [](const std::string &a, auto b) {
                            return a + " " + std::to_string(b.second);
                          });

      return errorMessage;
    }
  }

  return std::nullopt;
}

LogicalResult AIEX::NpuDmaMemcpyNdOp::verify() {
  BaseMemRefType buffer = getMemref().getType();
  const auto &targetModel = AIE::getTargetModel(*this);
  auto addressGranularity = targetModel.getAddressGenGranularity();

  if (getElementTypeBitwidth() > addressGranularity) {
    return emitOpError("Maximum element bit width allowed is ")
           << addressGranularity << "bits. ";
  }
  if (buffer.hasStaticShape() &&
      (buffer.getNumElements() * getElementTypeBitwidth()) <
          addressGranularity) {
    return emitOpError("Minimum data transfer size required is ")
           << addressGranularity << "bits. ";
  }
  if (!llvm::all_of(getMixedStrides(), [](OpFoldResult s) {
        return getConstantIntValue(s).has_value();
      }))
    return emitOpError("Only constant strides currently supported.");
  if (!llvm::all_of(getMixedSizes(), [](OpFoldResult s) {
        return getConstantIntValue(s).has_value();
      }))
    return emitOpError("Only constant sizes currently supported.");
  if (!llvm::all_of(getMixedOffsets(), [](OpFoldResult s) {
        return getConstantIntValue(s).has_value();
      }))
    return emitOpError("Only constant offsets currently supported.");

  llvm::SmallVector<int64_t, 4> inputSizes =
      llvm::map_to_vector(llvm::reverse(getMixedSizes()), [](OpFoldResult s) {
        return getConstantIntValue(s).value();
      });
  llvm::SmallVector<int64_t, 4> inputStrides =
      llvm::map_to_vector(llvm::reverse(getMixedStrides()), [](OpFoldResult s) {
        return getConstantIntValue(s).value();
      });
  llvm::SmallVector<int64_t, 4> hardwareSizes(4);
  llvm::SmallVector<int64_t, 4> hardwareStrides(4);
  getHardwareStridesWraps(targetModel, getOperation(), buffer, inputSizes,
                          inputStrides, hardwareSizes, hardwareStrides);
  int64_t offset = getOffsetInBytes();

  auto errorMessage = checkBurstLength(targetModel, getBurstLength());
  if (errorMessage.has_value()) {
    return emitOpError(errorMessage.value());
  }

  // The experimental HSA target uses this op on AIE1, skip all the AIE2
  // specific checks
  if (targetModel.getTargetArch() == AIE::AIEArch::AIE1)
    return success();

  if (offset % 4 != 0) {
    return emitOpError("Offset must be 4-byte-aligned.");
  }

  // dma_memcpy_nd transfers of the form [1, 1, 1, len][0, 0, 0, 1] do not
  // specify any data layout transformation, but simply express a contiguous
  // transfer of `len`. For backwards compatibility, we allow this to proceed
  // even if it exceeds the maximum stride/wrap size of any one dimension,
  // and simply do not lower any data layout transformations, since there is
  // no other way to express this at the dma_memcpy_nd interface otherwise.
  AIE::DeviceOp dev = getOperation()->getParentOfType<AIE::DeviceOp>();
  if (auto allocOp = AIE::ShimDMAAllocationOp::getForSymbol(
          dev, getMetadata().getRootReference())) {
    AIE::TileOp tile = allocOp.getTileOp();
    if (!tile) {
      return emitOpError("shim DMA allocation must reference a valid TileOp");
    }
    int col = tile.getCol();
    int row = tile.getRow();
    // A contiguous row-major ND access is also exempt from the ND wrap-size
    // limit: aie-dma-to-npu lowers it to linear mode (d0_size=d1_size=0),
    // and LinearizeContiguousTransfer canonicalizes it to explicit linear form.
    bool skipTransformationChecks =
        isLinearTransferWithoutTransformation() ||
        (targetModel.isShimNOCTile(col, row) &&
         AIEX::isContiguousTransfer(inputSizes, inputStrides));
    if (failed(verifyStridesWraps(*this, buffer, col, row, inputSizes,
                                  inputStrides, hardwareSizes, hardwareStrides,
                                  skipTransformationChecks))) {
      return failure();
    }
  }

  // packet header
  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.");
  }

  return success();
}

//===----------------------------------------------------------------------===//
// NpuDmaWaitOp
//===----------------------------------------------------------------------===//

LogicalResult AIEX::NpuDmaWaitOp::verify() {
  AIE::DeviceOp dev = (*this)->getParentOfType<AIE::DeviceOp>();
  // Some passes (e.g. aie-standard-lowering) use aiex ops outside a DeviceOp,
  // so we can't expect the device to always exist.
  if (dev && !dev.lookupSymbol(getSymbol()))
    return emitOpError("couldn't find symbol in parent device");
  return success();
}

//===----------------------------------------------------------------------===//
// NpuPushQueueOp
//===----------------------------------------------------------------------===//

LogicalResult AIEX::NpuPushQueueOp::verify() {
  const auto &targetModel = AIE::getTargetModel(*this);
  auto numBds = targetModel.getNumBDs(getColumn(), getRow());
  if (getBdId() > numBds)
    return emitOpError("BD ID exceeds the maximum ID.");
  if (getRepeatCount() > 255)
    return emitOpError("Repeat count exceeds the [0:255] range.");
  return success();
}

//===----------------------------------------------------------------------===//
// NpuWriteBdOp
//===----------------------------------------------------------------------===//

LogicalResult AIEX::NpuWriteBdOp::verify() {
  const auto &targetModel = AIE::getTargetModel(*this);
  auto numBds = targetModel.getNumBDs(getColumn(), getRow());
  bool isLinearTransfer =
      (getD0Size() >= 1) && (getD1Size() == 1) && (getIterationSize() == 0);
  if (getBdId() > numBds)
    return emitOpError("BD ID exceeds the maximum ID.");
  if (getPacketId() > 31)
    return emitOpError("Packet ID exceeds the maximum supported by 5 bits.");
  if (getPacketType() > 7)
    return emitOpError("Packet Type exceeds the maximum supported by 3 bits.");
  if (!isLinearTransfer && getD0Size() > 0x3FF)
    return emitOpError("D0 Size exceeds the [0:1023] range.");
  if (getD0Stride() > 0xFFFFF)
    return emitOpError("D0 Stride exceeds the [0:1M-1] range.");
  if (getD1Size() > 0x3FF)
    return emitOpError("D1 Size exceeds the [0:1023] range.");
  if (getD1Stride() > 0xFFFFF)
    return emitOpError("D1 Stride exceeds the [0:1M-1] range.");
  if (getD2Stride() > 0xFFFFF)
    return emitOpError("D2 Stride exceeds the [0:1M-1] range.");
  if (getIterationSize() > 0x3F)
    return emitOpError("Iteration Size exceeds the [0:63] range.");
  if (getIterationStride() > 0xFFFFF)
    return emitOpError("Iteration Stride exceeds the [0:1M-1] range.");
  if (targetModel.isShimNOCTile(getColumn(), getRow()) && getD2Size() != 0)
    return emitOpError("ShimTile only supports 3 dimensions of sizes.");
  if (targetModel.isShimNOCTile(getColumn(), getRow()) &&
      (getD0ZeroBefore() != 0 || getD0ZeroAfter() != 0 ||
       getD1ZeroBefore() != 0 || getD1ZeroAfter() != 0 ||
       getD2ZeroBefore() != 0 || getD2ZeroAfter() != 0))
    return emitOpError("ShimTile doesn't support zero padding.");
  if (!targetModel.isShimNOCTile(getColumn(), getRow()) &&
      getBurstLength() != 0)
    return emitOpError("Only ShimTiles support burst length.");
  auto errorMessage = checkBurstLength(targetModel, getBurstLength());
  if (errorMessage.has_value()) {
    return emitOpError(errorMessage.value());
  }

  return success();
}

//===----------------------------------------------------------------------===//
// NpuWrite32Op
//===----------------------------------------------------------------------===//

template <typename T>
static std::optional<uint32_t> getAbsoluteAddress(T *op) {
  AIE::DeviceOp device =
      op->getOperation()->template getParentOfType<AIE::DeviceOp>();
  if (!device) {
    op->emitError("Must be inside a device.");
    return std::nullopt;
  }
  const AIE::AIETargetModel &tm = device.getTargetModel();

  uint32_t address = 0;

  // If blockwrite references a buffer, the given address is understood to be
  // relative to the buffer's start address.
  if (op->getBuffer()) {
    AIE::BufferOp buffer = device.lookupSymbol<AIE::BufferOp>(*op->getBuffer());
    if (!buffer) {
      op->emitError() << "buffer '" << *op->getBuffer()
                      << "' not found in device";
      return std::nullopt;
    }

    if (!buffer.getAddress()) {
      mlir::InFlightDiagnostic err =
          op->emitError("referenced buffer must have address assigned");
      err.attachNote(buffer.getLoc()) << "This buffer must have an address.";
      return std::nullopt;
    }

    uint32_t col = buffer.getTileOp().getCol();
    uint32_t row = buffer.getTileOp().getRow();
    address = static_cast<uint32_t>(*buffer.getAddress()) +
              op->getAddress() * sizeof(uint32_t);
    address = ((col & 0xff) << tm.getColumnShift()) |
              ((row & 0xff) << tm.getRowShift()) | (address & 0xfffff);
  } else { // otherwise, the given address is absolute
    address = op->getAddress();
    std::optional<uint32_t> col = op->getColumn();
    std::optional<uint32_t> row = op->getRow();
    if (col && row) {
      // If col and row are set, only the lower 20 bits of the address are
      // used, and col and row dictate the upper bits (ignored)
      address = ((*col & 0xff) << tm.getColumnShift()) |
                ((*row & 0xff) << tm.getRowShift()) | (address & 0xfffff);
    }
  }

  return address;
}

std::optional<uint32_t> AIEX::NpuWrite32Op::getAbsoluteAddress() {
  return ::getAbsoluteAddress(this);
}

//===----------------------------------------------------------------------===//
// NpuUpdateFromScratchpadOp
//===----------------------------------------------------------------------===//

std::optional<uint32_t> AIEX::NpuUpdateFromScratchpadOp::getAbsoluteAddress() {
  return ::getAbsoluteAddress(this);
}

LogicalResult AIEX::NpuUpdateFromScratchpadOp::verify() {
  // StateTable has at most 32 entries (32-bit words).
  constexpr uint32_t kMaxStateTableEntries = 32;
  if (getStateTableIdx() >= kMaxStateTableEntries)
    return emitOpError("state_table_idx ")
           << static_cast<uint32_t>(getStateTableIdx())
           << " exceeds maximum StateTable index ("
           << (kMaxStateTableEntries - 1) << ").";

  // Cross-check against any npu.create_scratchpad ops in the same block: the
  // index must fit within the allocated scratchpad (size in 32-bit words).
  Block *block = (*this)->getBlock();
  if (block) {
    for (auto createOp : block->getOps<AIEX::NpuCreateScratchpadOp>()) {
      uint32_t sizeBytes = createOp.getSize();
      uint32_t numEntries = sizeBytes / 4;
      if (getStateTableIdx() >= numEntries) {
        return emitOpError("state_table_idx ")
               << static_cast<uint32_t>(getStateTableIdx())
               << " is out of bounds for scratchpad of size " << sizeBytes
               << " bytes (" << numEntries << " entries) created by "
               << createOp->getName() << ".";
      }
    }
  }
  return success();
}

//===----------------------------------------------------------------------===//
// NpuCreateScratchpadOp
//===----------------------------------------------------------------------===//

LogicalResult AIEX::NpuCreateScratchpadOp::verify() {
  // Only usage_type == 0 is currently supported by firmware.
  if (getUsageType() != 0)
    return emitOpError("usage_type must be 0 (got ")
           << static_cast<uint32_t>(getUsageType())
           << "); other layouts are not supported.";

  // The StateTable layout (usage_type 0) has a max of 32 32-bit-word entries,
  // i.e. a maximum total scratchpad size of 128 bytes.
  constexpr uint32_t kMaxScratchpadSizeBytes = 128;
  if (getSize() == 0)
    return emitOpError("size must be greater than 0.");
  if (getSize() % 4 != 0)
    return emitOpError("size (")
           << getSize() << ") must be a multiple of 4 bytes.";
  if (getSize() > kMaxScratchpadSizeBytes)
    return emitOpError("size (")
           << getSize() << " bytes) exceeds maximum scratchpad size of "
           << kMaxScratchpadSizeBytes << " bytes.";
  return success();
}

//===----------------------------------------------------------------------===//
// NpuMaskWrite32Op
//===----------------------------------------------------------------------===//

std::optional<uint32_t> AIEX::NpuMaskWrite32Op::getAbsoluteAddress() {
  return ::getAbsoluteAddress(this);
}

//===----------------------------------------------------------------------===//
// NpuBlockWriteOp
//===----------------------------------------------------------------------===//

std::optional<uint32_t> AIEX::NpuBlockWriteOp::getAbsoluteAddress() {
  return ::getAbsoluteAddress(this);
}

DenseIntElementsAttr AIEX::NpuBlockWriteOp::getDataWords() {
  Value memref = this->getData();
  DataLayout dataLayout = DataLayout::closest(*this);
  int64_t width = dataLayout.getTypeSizeInBits(
      cast<MemRefType>(memref.getType()).getElementType());
  if (width != 32) {
    emitWarning("Only 32-bit data type is supported for now");
    return nullptr;
  }

  memref::GetGlobalOp getGlobal = memref.getDefiningOp<memref::GetGlobalOp>();
  if (!getGlobal) {
    emitError("Only MemRefs from memref.get_global are supported");
    return nullptr;
  }

  auto global = dyn_cast_if_present<memref::GlobalOp>(
      (*this)->getParentOfType<AIE::DeviceOp>().lookupSymbol(
          getGlobal.getName()));
  if (!global) {
    emitError("Global symbol not found");
    return nullptr;
  }

  auto initVal = global.getInitialValue();
  if (!initVal) {
    emitError("Global symbol has no initial value");
    return nullptr;
  }

  auto data = dyn_cast<DenseIntElementsAttr>(*initVal);
  if (!data) {
    emitError("Global symbol initial value is not a dense int array");
    return nullptr;
  }

  return data;
}

//===----------------------------------------------------------------------===//
// DMAConfigureTaskOp
//===----------------------------------------------------------------------===//

std::optional<uint32_t> AIEX::DMAConfigureTaskOp::getFirstBdId() {
  Region &body = getBody();
  if (body.empty()) {
    return std::nullopt;
  }
  auto bd_ops = body.front().getOps<AIE::DMABDOp>();
  if (bd_ops.empty() && body.front().getNumSuccessors() == 1) {
    // Allow the first block to be empty and point to the entry point of the
    // chain. This allows for specifying cyclying BD chains (infinite loops)
    // within the constraints of MLIR syntax.
    Block &chain_entry = *body.front().getSuccessor(0);
    bd_ops = chain_entry.getOps<AIE::DMABDOp>();
  }
  if (bd_ops.empty()) {
    return std::nullopt;
  }
  AIE::DMABDOp bd = *bd_ops.begin();
  if (!bd.getBdId().has_value()) {
    return std::nullopt;
  }
  return bd.getBdId().value();
}

LogicalResult
AIEX::DMAConfigureTaskOp::canonicalize(AIEX::DMAConfigureTaskOp op,
                                       PatternRewriter &rewriter) {
  // Remove blocks that contain nothing but a terminator
  Region &body = op.getBody();
  bool did_rewrite = false;
  for (auto it = body.begin(); it != body.end(); ++it) {
    Block &block = *it;
    if (block.empty()) {
      continue;
    }
    auto ops_it = block.without_terminator();
    if (std::distance(ops_it.begin(), ops_it.end()) == 0) {
      rewriter.eraseOp(block.getTerminator());
      did_rewrite = true;
    }
  }
  if (did_rewrite) {
    return success();
  }
  return failure();
}

LogicalResult AIEX::DMAConfigureTaskOp::verify() {
  Region &body = getBody();
  for (auto it = body.begin(); it != body.end(); ++it) {
    Block &block = *it;
    if (block.empty()) {
      continue;
    }
    if (block.hasNoPredecessors() && !block.isEntryBlock()) {
      auto error = block.getTerminator()->emitError(
          "Block ending in this terminator does not form a chain with "
          "entry block.");
      return failure();
    }

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

    // This is a layering violation on the DMABDOps, but they are never verified
    // otherwise Because DMAConfigureTaskOps are not yet merged into the AIE
    // dialect. The normal DMABDOp verify operation will skip over any BD inside
    // a DMAConfigureTaskOp
    LogicalResult result = success();
    block.walk([&](AIE::DMABDOp bd) {
      if (bd.getBurstLength() != 0 &&
          !targetModel.isShimNOCTile(getTileID().col, getTileID().row)) {
        bd.emitOpError("Burst length is only supported in Shim NOC tiles that "
                       "are connected to the memory-mapped NOC.");
        result = failure();
      }
    });
    if (failed(result)) {
      return result;
    }
  }
  return success();
}

//===----------------------------------------------------------------------===//
// DMAStartBdChainOp
//===----------------------------------------------------------------------===//

AIE::BDChainOp AIEX::DMAStartBdChainOp::getBDChainOp() {
  AIE::DeviceOp device = (*this)->getParentOfType<AIE::DeviceOp>();
  AIE::BDChainOp chain = device.lookupSymbol<AIE::BDChainOp>(getSymbol());
  return chain;
}

LogicalResult AIEX::DMAStartBdChainOp::verify() {
  AIE::BDChainOp chain = getBDChainOp();
  if (!chain) {
    return emitOpError("symbol does not reference valid BD chain");
  }

  auto actualArgTypes = getArgs().getTypes();
  auto expectedArgTypes = chain.getRegion().getArgumentTypes();
  if (actualArgTypes.size() != expectedArgTypes.size()) {
    return emitOpError("Number of arguments mismatches.");
  }
  for (unsigned i = 0, n = expectedArgTypes.size(); i < n; i++) {
    if (actualArgTypes[i] != expectedArgTypes[i]) {
      return emitOpError("Argument ") << (i + 1) << " types mismatch: "
                                      << "expected " << expectedArgTypes[i]
                                      << " but got " << actualArgTypes[i];
    }
  }
  return success();
}

//===----------------------------------------------------------------------===//
// NpuControlPacketOp
//===----------------------------------------------------------------------===//

uint32_t AIEX::NpuControlPacketOp::getRowFromAddr() {
  const auto &targetModel = AIE::getTargetModel(*this);
  uint32_t addr = getAddress();
  uint32_t rowInt = (addr >> targetModel.getRowShift()) & 0x1f;
  return rowInt;
}

uint32_t AIEX::NpuControlPacketOp::getColumnFromAddr() {
  const auto &targetModel = AIE::getTargetModel(*this);
  uint32_t addr = getAddress();
  uint32_t colInt = (addr >> targetModel.getColumnShift()) & 0x1f;
  return colInt;
}

//===----------------------------------------------------------------------===//
// SetLockOp
//===----------------------------------------------------------------------===//

LogicalResult AIEX::SetLockOp::verify() {
  const auto &targetModel = AIE::getTargetModel(*this);

  if (targetModel.getTargetArch() == AIE::AIEArch::AIE1)
    return emitOpError("SetLockOp is not supported on AIE1.");

  if (getValue() > targetModel.getMaxLockValue())
    return emitOpError("Lock value exceeds the maximum value of " +
                       std::to_string(targetModel.getMaxLockValue()));

  auto lockOp = getLockOp();
  auto lockIDOpt = getLockOp().getLockID();
  // Note that the lockID may not be assigned initially, so lets wait until it
  // is to verify the lockID dependent conditions
  if (!lockIDOpt) {
    return success();
  }

  auto col = lockOp.colIndex();
  auto row = lockOp.rowIndex();
  uint32_t lockID = lockOp.getLockIDValue();

  if (lockID >= targetModel.getNumLocks(col, row)) {
    return emitOpError("Lock ID out of range for given tile. Max ID: " +
                       std::to_string(targetModel.getNumLocks(col, row) - 1));
  }

  if (!targetModel.getLocalLockAddress(lockID, lockOp.getTileID())) {
    return emitOpError("Invalid lock ID and tile combination when trying to "
                       "retrieve the local lock address.");
  }

  return success();
}

//===----------------------------------------------------------------------===//
// BlockFloatingPointType
//===----------------------------------------------------------------------===//
uint64_t AIEX::BlockFloatType::getTotalSizeInBits() const {
  return getBlockSize() * getMantissaBits() + getExponentBits() +
         getSubtileShiftBits();
}

llvm::TypeSize AIEX::BlockFloatType::getTypeSizeInBits(
    const mlir::DataLayout &dataLayout,
    mlir::DataLayoutEntryListRef params) const {
  return llvm::TypeSize::getFixed(getTotalSizeInBits());
}

uint64_t AIEX::BlockFloatType::getABIAlignment(
    const mlir::DataLayout &dataLayout,
    mlir::DataLayoutEntryListRef params) const {
  // For the purposes of the data movement operations, we want all types to be
  // packed <=> ABI alignment is 1.
  return 1;
}

std::optional<AIEX::BlockFloatType::BlockFormat>
AIEX::BlockFloatType::getBlockFormat(StringRef blockType) {
  static const llvm::StringMap<AIEX::BlockFloatType::BlockFormat>
      blockFormatsMap = {
          {"v8bfp16ebs8", {8, 8, 8, 0}},
          {"v16bfp16ebs16", {16, 8, 8, 0}},
      };

  auto it = blockFormatsMap.find(blockType);
  if (it != blockFormatsMap.end()) {
    return it->second;
  }

  return std::nullopt;
}

LogicalResult
AIEX::BlockFloatType::verify(function_ref<InFlightDiagnostic()> emitError,
                             StringRef block_type) {
  if (!getBlockFormat(block_type))
    return emitError() << "Invalid block type: " << block_type
                       << ". Known types are: v8bfp16ebs8, v16bfp16ebs16.";

  return success();
}

//===----------------------------------------------------------------------===//
// ConfigureOp
//===----------------------------------------------------------------------===//

AIE::DeviceOp AIEX::ConfigureOp::getReferencedDeviceOp() {
  ModuleOp moduleOp = this->getOperation()->getParentOfType<ModuleOp>();
  if (!moduleOp) {
    emitError("aiex.configure must be inside of a module");
    return nullptr;
  }
  Operation *maybeReferencedDevice =
      SymbolTable::lookupSymbolIn(moduleOp.getOperation(), getSymbolAttr());
  if (!maybeReferencedDevice) {
    emitError("No such device: '") << getSymbolAttr() << "'";
    return nullptr;
  }
  AIE::DeviceOp referencedDevice =
      llvm::dyn_cast<AIE::DeviceOp>(maybeReferencedDevice);
  if (!referencedDevice) {
    emitError("Not a device: '") << getSymbolAttr() << "'";
    return nullptr;
  }
  return referencedDevice;
}

//===----------------------------------------------------------------------===//
// ReadParameterOp
//===----------------------------------------------------------------------===//

LogicalResult AIEX::ReadParameterOp::verify() {
  auto device = (*this)->getParentOfType<AIE::DeviceOp>();
  if (!device) {
    return emitOpError("must be inside an aie.device");
  }
  if (!(*this)->getParentOfType<AIE::CoreOp>()) {
    return emitOpError("must be inside an aie.core");
  }
  auto module = (*this)->getParentOfType<ModuleOp>();
  if (!module || !module.lookupSymbol<AIEX::ParameterOp>(getParameter())) {
    return emitOpError("references unknown parameter '")
           << getParameter()
           << "' (aiex.parameter ops are declared at module scope)";
  }
  return success();
}

LogicalResult AIEX::ConfigureOp::verify() {
  AIE::DeviceOp parentDev = getOperation()->getParentOfType<AIE::DeviceOp>();
  AIE::DeviceOp referencedDev = getReferencedDeviceOp();
  if (!referencedDev) {
    return failure();
  }
  if (parentDev.getDevice() != referencedDev.getDevice()) {
    emitError("Device types do not match: '")
        << AIE::stringifyAIEDevice(parentDev.getDevice()) << "' vs. '"
        << AIE::stringifyAIEDevice(referencedDev.getDevice()) << "'";
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// RunOp
//===----------------------------------------------------------------------===//

AIE::DeviceOp AIEX::RunOp::getCalleeDeviceOp() {
  AIEX::ConfigureOp configureOp =
      getOperation()->getParentOfType<AIEX::ConfigureOp>();
  if (!configureOp) {
    return nullptr;
  }
  AIE::DeviceOp referencedDevice = configureOp.getReferencedDeviceOp();
  return referencedDevice;
}

AIE::RuntimeSequenceOp AIEX::RunOp::getCalleeRuntimeSequenceOp() {
  AIEX::ConfigureOp configureOp =
      getOperation()->getParentOfType<AIEX::ConfigureOp>();
  if (!configureOp) {
    return nullptr;
  }
  AIE::DeviceOp referencedDevice = configureOp.getReferencedDeviceOp();
  if (!referencedDevice) {
    return nullptr;
  }

  Operation *maybeRuntimeSequence =
      SymbolTable::lookupSymbolIn(referencedDevice, getRuntimeSequenceSymbol());

  if (!maybeRuntimeSequence) {
    return nullptr;
  }
  AIE::RuntimeSequenceOp runtimeSequence =
      llvm::dyn_cast<AIE::RuntimeSequenceOp>(maybeRuntimeSequence);
  if (!runtimeSequence) {
    return nullptr;
  }

  return runtimeSequence;
}

//===----------------------------------------------------------------------===//
// NpuLoadPdiOp
//===----------------------------------------------------------------------===//

LogicalResult AIEX::NpuLoadPdiOp::canonicalize(AIEX::NpuLoadPdiOp op,
                                               PatternRewriter &rewriter) {
  // Check for back-to-back identical load_pdi ops and remove duplicates
  Operation *nextOp = op->getNextNode();
  if (!nextOp)
    return failure();

  // Check if next op is also a NpuLoadPdiOp
  auto nextLoadPdi = dyn_cast<AIEX::NpuLoadPdiOp>(nextOp);
  if (!nextLoadPdi)
    return failure();

  // Check if they are identical (all attributes match)
  if (op.getDeviceRefAttr() == nextLoadPdi.getDeviceRefAttr() &&
      op.getId() == nextLoadPdi.getId() &&
      op.getSize() == nextLoadPdi.getSize() &&
      op.getAddress() == nextLoadPdi.getAddress()) {
    // Erase the first one, keeping the second
    rewriter.eraseOp(op);
    return success();
  }

  return failure();
}