AIEXDialect.cpp

Go to the documentation of this file.

//===- AIEXDialect.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/AIEX/IR/AIEXDialect.h"
#include "aie/Dialect/AIEX/Utils/BdLowering.h"
#include "aie/Dialect/AIEX/Utils/DmaDecomposition.h"
 
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/DialectImplementation.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/RegionGraphTraits.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/SCCIterator.h"
#include "llvm/ADT/STLExtras.h"
 
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/CheckedArithmetic.h"
#include "llvm/Support/MathExtras.h"
#include "llvm/Support/TypeSize.h"
 
#include <cstdint>
#include <limits>
#include <numeric>
 
using namespace mlir;
using namespace xilinx;
 
#include "aie/Dialect/AIEX/IR/AIEXDialect.cpp.inc"
 
#include "aie/Dialect/AIEX/IR/AIEXEnums.cpp.inc"
 
#include "aie/Dialect/AIEX/IR/AIEXInterfaces.cpp.inc"
 
#define GET_TYPEDEF_CLASSES
#include "aie/Dialect/AIEX/IR/AIEXTypes.cpp.inc"
 
#define GET_ATTRDEF_CLASSES
#include "aie/Dialect/AIEX/IR/AIEXAttrDefs.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"
      >();
  addAttributes<
#define GET_ATTRDEF_LIST
#include "aie/Dialect/AIEX/IR/AIEXAttrDefs.cpp.inc"
      >();
}
 
} // namespace xilinx::AIEX
 
using xilinx::AIE::parseTypedDynamicIndexList;
using xilinx::AIE::printTypedDynamicIndexList;
 
#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
  llvm::fill(sizes, 0);
  llvm::fill(strides, 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;
  }
 
  ConstStridePolicy policy;
  int64_t inS[4] = {inputSizes[0], inputSizes[1], inputSizes[2], inputSizes[3]};
  int64_t inT[4] = {inputStrides[0], inputStrides[1], inputStrides[2],
                    inputStrides[3]};
  int64_t outS[4], outT[4];
  encodeHardwareStridesWraps(policy, elemWidth, addressGranularity, inS, inT,
                             outS, outT);
  for (int i = 0; i < 4; i++) {
    sizes[i] = outS[i];
    strides[i] = outT[i];
  }
}
 
mlir::LogicalResult AIEX::verifyStridesWraps(
    mlir::Operation *forOp,
    llvm::function_ref<mlir::InFlightDiagnostic()> emitError,
    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());
 
  // ShimPLTiles have no ShimDMA, so BD field widths are meaningless here.
  if (!targetModel.isCoreTile(tileCol, tileRow) &&
      !targetModel.isMemTile(tileCol, tileRow) &&
      !targetModel.isShimNOCTile(tileCol, tileRow))
    return emitError() << "Unsupported tile type at (" << tileCol << ", "
                       << tileRow << ") Must be ShimNOC, Mem or Core.";
 
  uint32_t wrap_bits = targetModel.getDmaBdWrapBits(tileCol, tileRow);
  uint32_t step_bits = targetModel.getDmaBdStepBits(tileCol, tileRow);
 
  for (int i = 0; i < 4; i++) {
    if (inputSizes[i] <= 0) {
      return emitError() << "Size " << i << " must be a positive integer.";
    }
  }
 
  if (!isConstMultipleOfGranule(inputSizes[0], elemWidth, addressGranularity)) {
    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 emitError() << 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 emitError() << "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 emitError() << "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 (!isConstMultipleOfGranule(inputStrides[i], elemWidth,
                                  addressGranularity)) {
      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 emitError() << msg.str();
    }
  }
  // The innermost dimension steps whole granules, so for elements of another
  // width it can only be contiguous.
  if (elemWidth != addressGranularity && inputSizes[0] > 1 &&
      inputStrides[0] != 1)
    return emitError() << "Stride 0 is " << inputStrides[0]
                       << " elements, but must be 1 for " << (elemWidth / 8)
                       << "-byte elements: the DMA moves whole "
                       << (addressGranularity / 8) << "-byte words.";
 
  auto outOfRange = [&](const char *what, int dim, int64_t written,
                        int64_t encoded, int64_t lo, int64_t hi) {
    return emitError() << what << " " << dim << " is " << written
                       << " (encoded as " << encoded << "), which exceeds the ["
                       << lo << ":" << hi << "] range.";
  };
 
  int64_t maxWrap = (1 << wrap_bits) - 1;
  int64_t maxStep = 1 << step_bits;
 
  for (int dim : {0, 1}) {
    if (!skipTransformationChecks && hardwareSizes[dim] > maxWrap)
      return outOfRange("Size", dim, inputSizes[dim], hardwareSizes[dim], 0,
                        maxWrap);
  }
  // A zero d3 stride is a pure repeat: the lowerings leave the iteration fields
  // 0 and carry the count on the queue push, so the repeat limit applies.
  bool pureRepeat = inputStrides[3] == 0;
  int64_t maxCount =
      pureRepeat ? targetModel.getMaxRepeatCount() + 1
                 : 1LL << targetModel.getDmaBdIterBits(tileCol, tileRow);
  if (inputSizes[3] > maxCount)
    return emitError() << (pureRepeat ? "repeat count " : "iteration count ")
                       << inputSizes[3] << " exceeds the [1:" << maxCount
                       << "] range.";
  for (int dim : {0, 1, 2}) {
    if (hardwareStrides[dim] > maxStep - 1)
      return outOfRange("Stride", dim, inputStrides[dim], hardwareStrides[dim],
                        1, maxStep);
  }
  // strides[3] exceeding the range is ok iff the sizes[3] is one, which is
  // checked below
  if (hardwareStrides[3] > maxStep - 1 && hardwareSizes[3] > 0)
    return outOfRange("Stride", 3, inputStrides[3], hardwareStrides[3], 1,
                      maxStep);
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// UseTokenOp
//===----------------------------------------------------------------------===//
 
// A field flagged as a register must name one of r0..r23. The bound cannot be
// an attribute constraint, because the same field holds a full 32-bit
// immediate when unflagged.
LogicalResult AIEX::CertWrite32DOp::verify() {
  if (getAddressIsReg() && getAddress() > 23)
    return emitOpError("has address_is_reg, so address ")
           << getAddress() << " must be a register index in [0, 23]";
  if (getValueIsReg() && getValue() > 23)
    return emitOpError("has value_is_reg, so value ")
           << getValue() << " must be a register index in [0, 23]";
  return success();
}
 
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
 */
uint32_t AIEX::NpuDmaMemcpyNdOp::getAxcacheOrDefault() {
  return getAxcache().value_or(AIE::getTargetModel(*this).getDefaultAxCache());
}
 
int64_t AIEX::NpuDmaMemcpyNdOp::getOffsetInBytes() {
  llvm::SmallVector<int64_t, 4> offsets =
      llvm::map_to_vector(llvm::reverse(getMixedOffsets()), [](OpFoldResult s) {
        return getConstantIntValue(s).value();
      });
  auto strides = llvm::to_vector<4>(llvm::reverse(getMixedStrides()));
  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;
  // A dimension only contributes to the byte offset when its (constant) offset
  // is non-zero; a runtime stride paired with a zero offset is fine and must
  // not be forced to a constant. The verifier requires any stride multiplied by
  // a non-zero offset to be constant.
  for (size_t i = 0; i < R; i++) {
    if (offsets[i] == 0)
      continue;
    auto strideConst = getConstantIntValue(strides[i]);
    assert(strideConst &&
           "verifier requires a stride paired with a non-zero offset to be "
           "constant");
    offset += offsets[i] * (*strideConst) * 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;
}
 
bool AIEX::isLinearModeTransfer(const AIE::AIETargetModel &targetModel, int col,
                                int row, llvm::ArrayRef<int64_t> sizes,
                                llvm::ArrayRef<int64_t> strides) {
  return isLinearTransfer(sizes, strides) ||
         (targetModel.isShimNOCTile(col, row) &&
          isContiguousTransfer(sizes, strides));
}
 
// 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, newOffsetsOuter, newSizesOuter, newStridesOuter);
    return mlir::success();
  }
};
} // namespace
 
void AIEX::NpuDmaMemcpyNdOp::getCanonicalizationPatterns(
    mlir::RewritePatternSet &results, mlir::MLIRContext *context) {
  results.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 = llvm::find_if(bel, [=](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;
}
 
// Verify the supported scope for a dma_memcpy_nd carrying runtime (SSA)
// offsets/sizes/strides, hard-erroring on any statically-provable violation.
// The scope here is exactly what the dynamic BD encoder (AIEDmaToNpu.cpp
// lowerDynamic) can lower; runtime values are never silently masked.
LogicalResult AIEX::NpuDmaMemcpyNdOp::verifyDynamicSizesStrides(
    const AIE::AIETargetModel &targetModel, mlir::BaseMemRefType buffer,
    AIE::ShimDMAAllocationOp allocOp) {
  // Shim NOC only.
  if (!allocOp)
    return emitOpError(
        "runtime sizes/strides require a shim_dma_allocation to resolve the "
        "tile; none found.");
  AIE::TileOp tile = allocOp.getTileOp();
  // Coordinate-dependent checks must wait until logical tiles are placed.
  if (!tile)
    return success();
  if (!targetModel.isShimNOCTile(tile.getCol(), tile.getRow()))
    return emitOpError(
        "runtime sizes/strides are only supported for shim NOC tile DMAs.");
 
  // No zero-padding with runtime sizes/strides.
  if (getD0ZeroBefore() || getD1ZeroBefore() || getD2ZeroBefore() ||
      getD0ZeroAfter() || getD1ZeroAfter() || getD2ZeroAfter())
    return emitOpError(
        "zero padding is not supported with runtime sizes/strides.");
 
  // Per-field runtime values are allowed: any individual size/stride may be an
  // SSA value while others stay constant. Only the runtime-dependent fields
  // need runtime handling; the encoder produces the same word for a constant
  // operand either way, so the static ≡ dynamic byte-identity holds field by
  // field.
  auto sizes = getMixedSizes();
  auto strides = getMixedStrides();
 
  // The innermost stride may be runtime like any other dimension: the encoder
  // resolves its collapse-to-zero case with a select, and its realizability
  // (unit stride, or granule-aligned) is enforced below for constants and by a
  // host-side check (buildBdWords) for runtime values.
 
  // (The memref must also trace to a runtime-sequence block argument through
  // static subview/cast offsets; that structural check, with the same clean
  // diagnostic, is enforced by the lowering in AIEDmaToNpu.cpp, which already
  // owns the trace utility. It is not duplicated here to keep the dialect
  // verifier free of the analysis-layer dependency.)
 
  // A constant size must fit its hardware wrap field (an out-of-range constant
  // is a hard error, never a silent truncation); runtime sizes are guarded at
  // lowering. Bounds are checked on the ENCODED value, matching
  // verifyStridesWraps: d0 wrap scaled to granules (size * elemWidth / gran),
  // iteration wrap biased by -1, d1 a raw element count. Sizes outermost-first.
  DataLayout dataLayout = DataLayout::closest(getOperation());
  uint64_t elemWidth = dataLayout.getTypeSizeInBits(buffer.getElementType());
  uint32_t gran = targetModel.getAddressGenGranularity();
  llvm::SmallVector<mlir::OpFoldResult, 4> sizesRev(llvm::reverse(sizes));
  auto checkSize = [&](mlir::OpFoldResult ofr, int64_t hwVal, int64_t hi,
                       llvm::StringRef what) -> LogicalResult {
    if (!getConstantIntValue(ofr))
      return success();
    if (hwVal < 0 || hwVal > hi)
      return emitOpError(what) << " hardware value " << hwVal
                               << " exceeds hardware range [0:" << hi << "].";
    return success();
  };
  auto hwSize = [&](mlir::OpFoldResult ofr) {
    return getConstantIntValue(ofr).value_or(0);
  };
  int64_t d0Hw = (int64_t)(hwSize(sizesRev[0]) * elemWidth / gran);
  int64_t d1Hw = hwSize(sizesRev[1]);
  int64_t iterRaw = hwSize(sizesRev[3]);
  int64_t iterHw = iterRaw > 1 ? iterRaw - 1 : 0;
  std::optional<int64_t> outerStride = getConstantIntValue(strides.front());
  bool pureRepeat = outerStride && *outerStride == 0;
  AIE::AIETileType tileType =
      targetModel.getTileType(tile.getCol(), tile.getRow());
  int64_t wrapMax = (1LL << targetModel.getDmaBdWrapBits(tileType)) - 1;
  int64_t maxCount = pureRepeat ? targetModel.getMaxRepeatCount() + 1
                                : 1LL << targetModel.getDmaBdIterBits(tileType);
  if (failed(checkSize(sizesRev[0], d0Hw, wrapMax, "d0 size")) ||
      failed(checkSize(sizesRev[1], d1Hw, wrapMax, "d1 size")) ||
      failed(checkSize(sizesRev[3], iterHw, maxCount - 1,
                       pureRepeat ? "repeat count" : "iteration size")))
    return failure();
 
  // Realizability of the CONSTANT size/stride operands (divisibility +
  // positivity, innermost-first). Runtime operands get a host-side check at
  // lowering time. Shared with the dma_task path.
  llvm::SmallVector<mlir::OpFoldResult, 4> stridesRev(llvm::reverse(strides));
  if (failed(verifyConstBdRealizability(getOperation(), sizesRev, stridesRev,
                                        elemWidth, gran)))
    return failure();
 
  // Runtime fields are checked at dispatch instead; see buildBdWords.
 
  auto errorMessage = checkBurstLength(targetModel, getBurstLength());
  if (errorMessage.has_value())
    return emitOpError(errorMessage.value());
 
  return success();
}
 
LogicalResult AIEX::NpuDmaMemcpyNdOp::verify() {
  BaseMemRefType buffer = getMemref().getType();
  const auto &targetModel = AIE::getTargetModel(*this);
  auto addressGranularity = targetModel.getAddressGenGranularity();
 
  if (getOffsetParameterAttr() || getOffsetStateTableIdxAttr()) {
    uint64_t elemBitWidth = buffer.getElementTypeBitWidth();
    if (elemBitWidth == 0 || (elemBitWidth % 8) != 0)
      return emitOpError("offset_parameter requires a whole-byte element type");
  }
 
  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. ";
  }
  bool allStridesConstant = llvm::all_of(getMixedStrides(), [](OpFoldResult s) {
    return getConstantIntValue(s).has_value();
  });
  bool allSizesConstant = llvm::all_of(getMixedSizes(), [](OpFoldResult s) {
    return getConstantIntValue(s).has_value();
  });
  bool allOffsetsConstant = llvm::all_of(getMixedOffsets(), [](OpFoldResult s) {
    return getConstantIntValue(s).has_value();
  });
 
  bool hasLengthParameter =
      getLengthParameterAttr() || getLengthStateTableIdxAttr();
  // Runtime offsets, sizes and strides lower only through EmitC, which has no
  // scratchpad.
  if (hasLengthParameter &&
      (!allStridesConstant || !allSizesConstant || !allOffsetsConstant))
    return emitOpError(
        "length_parameter requires constant offsets, sizes and strides");
 
  // The dynamic scope check needs the allocation resolved by verifySymbolUses.
  if (!allStridesConstant || !allSizesConstant || !allOffsetsConstant)
    return success();
 
  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();
      });
  int64_t offset = getOffsetInBytes();
 
  auto errorMessage = checkBurstLength(targetModel, getBurstLength());
  if (errorMessage.has_value()) {
    return emitOpError(errorMessage.value());
  }
 
  if (hasLengthParameter) {
    if (failed(AIE::verifyLengthParameter(
            *this, getLengthUnit(), buffer,
            inputSizes[0] * inputSizes[1] * inputSizes[2],
            AIEX::isContiguousTransfer(inputSizes, inputStrides),
            llvm::ArrayRef<int64_t>(inputSizes).take_front(3))))
      return failure();
  }
 
  // AIE1 does not use the NPU alignment and packet-header checks.
  if (targetModel.getTargetArch() == AIE::AIEArch::AIE1)
    return success();
 
  if (offset % 4 != 0) {
    return emitOpError("Offset must be 4-byte-aligned.");
  }
 
  // packet header
  if (auto packetInfo = getPacket()) {
    if (failed(AIE::verifyAssignedPacket(*this, *packetInfo)))
      return failure();
    if (packetInfo->getPktType() > 7)
      return emitOpError("Packet type field can only hold 3 bits.");
    if (packetInfo->assignedId() > 31)
      return emitOpError("Packet ID field can only hold 5 bits.");
  }
 
  return success();
}
 
// The verifier-owned collection shares the device symbol table across memcpy
// checks without retaining allocation pointers between verification calls.
LogicalResult
AIEX::NpuDmaMemcpyNdOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  BaseMemRefType buffer = getMemref().getType();
  const auto &targetModel = AIE::getTargetModel(*this);
  bool allStridesConstant = llvm::all_of(getMixedStrides(), [](OpFoldResult s) {
    return getConstantIntValue(s).has_value();
  });
  bool allSizesConstant = llvm::all_of(getMixedSizes(), [](OpFoldResult s) {
    return getConstantIntValue(s).has_value();
  });
  bool allOffsetsConstant = llvm::all_of(getMixedOffsets(), [](OpFoldResult s) {
    return getConstantIntValue(s).has_value();
  });
  bool isDynamic =
      !allStridesConstant || !allSizesConstant || !allOffsetsConstant;
 
  bool hasLengthParameter =
      getLengthParameterAttr() || getLengthStateTableIdxAttr();
  bool isAIE1 = targetModel.getTargetArch() == AIE::AIEArch::AIE1;
  AIE::DeviceOp dev = getOperation()->getParentOfType<AIE::DeviceOp>();
  if (!dev && isAIE1 && !isDynamic && !hasLengthParameter)
    return success();
  auto allocOp = symbolTable.lookupSymbolIn<AIE::ShimDMAAllocationOp>(
      dev, getMetadata().getRootReference());
  if (isDynamic)
    return verifyDynamicSizesStrides(targetModel, buffer, allocOp);
 
  // Static transfers can precede allocation lowering, so defer tile checks
  // until metadata names a shim DMA allocation.
  if (!allocOp)
    return success();
  auto tile = allocOp.getTileOp();
  // Defer coordinate-dependent checks until logical tiles are placed.
  if (!tile) {
    if (isAIE1 && hasLengthParameter)
      return emitOpError("length_parameter is only supported on AIE2 and "
                         "AIE2P devices");
    return success();
  }
  int col = tile.getCol();
  int row = tile.getRow();
  if (hasLengthParameter &&
      failed(AIE::verifyLengthParameterTile(*this, targetModel, col, row)))
    return failure();
 
  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);
 
  // 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.
  // 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 = AIEX::isLinearModeTransfer(
      targetModel, col, row, inputSizes, inputStrides);
  // An oversized non-contiguous pattern that aie-decompose-large-dma-bd can
  // split into hardware-legal sub-transfers is also allowed to verify: the
  // pass rewrites it before BD lowering. Truly undecomposable patterns (e.g.
  // oversized strides) still fail below. isDecomposableNdDmaPattern
  // suppresses its own diagnostics, so no stray error is emitted for the
  // accepted case.
  llvm::SmallVector<int64_t, 4> inputOffsets =
      llvm::map_to_vector(llvm::reverse(getMixedOffsets()), [](OpFoldResult s) {
        return getConstantIntValue(s).value();
      });
  // aie-decompose-large-dma-bd leaves a transfer with a runtime length
  // alone, so it must be legal as written.
  // AIE1 patterns must fit directly, even when the decomposition pass is
  // scheduled: verification must not rely on it to legalize AIE1 transfers.
  if (!isAIE1 && !hasLengthParameter &&
      AIEX::isDecomposableNdDmaPattern(getOperation(), buffer, targetModel, col,
                                       row, inputOffsets, inputSizes,
                                       inputStrides)) {
    // ok: will be decomposed before lowering
  } else if (failed(verifyStridesWraps(
                 *this, [&] { return emitOpError(); }, buffer, col, row,
                 inputSizes, inputStrides, hardwareSizes, hardwareStrides,
                 skipTransformationChecks))) {
    return failure();
  }
 
  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());
  // bd_id and repeat_count are SSA operands; range-check them only when they
  // are compile-time constants. A runtime repeat_count is checked on the host
  // when lowered (AIEDmaToNpu); a runtime bd_id comes from the tile's BD pool.
  if (std::optional<uint32_t> bdId = getConstantIntOperand(getBdId());
      bdId && *bdId >= numBds)
    return emitOpError("BD ID exceeds the maximum ID.");
  uint32_t maxRepeat = targetModel.getMaxRepeatCount();
  if (std::optional<uint64_t> repeatCount =
          getConstantInt64Operand(getRepeatCount());
      repeatCount && *repeatCount > maxRepeat)
    return emitOpError("Repeat count exceeds the [0:")
           << maxRepeat << "] range.";
  if (getRepeatStateTableIdxAttr()) {
    AIE::AIEArch arch = targetModel.getTargetArch();
    if (arch != AIE::AIEArch::AIE2 && arch != AIE::AIEArch::AIE2p)
      return emitOpError("repeat_parameter is only supported on AIE2 and "
                         "AIE2P devices");
    if (!targetModel.isShimNOCTile(getColumn(), getRow()) &&
        !targetModel.isMemTile(getColumn(), getRow()))
      return emitOpError("repeat_parameter is only supported on shim NOC "
                         "and mem tiles, got tile (")
             << getColumn() << ", " << getRow() << ")";
    if (getConstantInt64Operand(getRepeatCount()) != uint64_t{0})
      return emitOpError("takes its repeat count from the scratchpad, so its "
                         "repeat_count operand must be 0");
    std::optional<uint32_t> bdId = getConstantIntOperand(getBdId());
    if (!bdId)
      return emitOpError("takes its repeat count from the scratchpad, so its "
                         "bd_id must be a constant");
    if (*bdId % 4 != 0)
      return emitOpError("takes its repeat count from the scratchpad, whose "
                         "update clears the low two bits of the queue word, "
                         "so its bd_id must be a multiple of 4, got ")
             << *bdId;
  }
  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.");
  int64_t oooId = getOutOfOrderId();
  if (oooId < 0 ||
      static_cast<uint64_t>(oooId) > targetModel.getMaxOutOfOrderId())
    return emitOpError("out_of_order_id must be in [0, ")
           << targetModel.getMaxOutOfOrderId() << "].";
  if (oooId != 0 && getEnablePacket() == 0)
    return emitOpError("out_of_order_id requires a packet-enabled BD");
 
  // Every value on this op is already the hardware-encoded field value (wrap
  // fields unbiased, stepsize/iteration fields biased actual-1), so the
  // legal encoded range for a B-bit field is simply [0, 2^B - 1].
  AIE::AIETileType tileType = targetModel.getTileType(getColumn(), getRow());
  uint32_t wrapBits = targetModel.getDmaBdWrapBits(tileType);
  uint32_t stepBits = targetModel.getDmaBdStepBits(tileType);
  uint32_t iterBits = targetModel.getDmaBdIterBits(tileType);
  uint32_t maxWrap = wrapBits > 0 ? (1u << wrapBits) - 1 : 0;
  uint32_t maxStep = stepBits > 0 ? (1u << stepBits) - 1 : 0;
  uint32_t maxIter = iterBits > 0 ? (1u << iterBits) - 1 : 0;
 
  if (!isLinearTransfer && getD0Size() > maxWrap)
    return emitOpError() << "D0 Size exceeds the [0:" << maxWrap << "] range.";
  if (getD0Stride() > maxStep)
    return emitOpError() << "D0 Stride exceeds the [0:" << maxStep
                         << "] range.";
  if (getD1Size() > maxWrap)
    return emitOpError() << "D1 Size exceeds the [0:" << maxWrap << "] range.";
  if (getD1Stride() > maxStep)
    return emitOpError() << "D1 Stride exceeds the [0:" << maxStep
                         << "] range.";
  if (getD2Stride() > maxStep)
    return emitOpError() << "D2 Stride exceeds the [0:" << maxStep
                         << "] range.";
  if (getIterationSize() > maxIter)
    return emitOpError() << "Iteration Size exceeds the [0:" << maxIter
                         << "] range.";
  if (static_cast<uint32_t>(getIterationStride()) > maxStep)
    return emitOpError() << "Iteration Stride exceeds the [0:" << maxStep
                         << "] range.";
  // buffer_length is the full 32-bit word on a shim NOC tile but only 17 bits
  // on a mem tile and 14 on a core tile, where the packer masks it.
  uint64_t maxLen = targetModel.getDmaBdMaxLen(tileType);
  if (getBufferLength() > maxLen)
    return emitOpError() << "Buffer length exceeds the [0:" << maxLen
                         << "] 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.");
 
  // Pad field widths (AIE2P ArchSpec Table 3-34, in 32-bit words): D0 is a
  // 6-bit field (max 63), D1 a 5-bit field (max 31), D2 a 4-bit field (max
  // 15).
  //
  // Placement is a deliberate compromise. DMABDOp::verify() would be the
  // better home for this check, since its diagnostic would point at the
  // user's own `aie.dma_bd`. But DMABDOp::verify() early-returns for BDs
  // nested inside `aiex.dma_configure_task`, and the only case with
  // hardware evidence for this defect is on that task path -- a check
  // placed only in DMABDOp::verify() would miss it. NpuWriteBdOp is where
  // both the static and task runtime paths converge, so putting it here
  // covers both, at the cost of the diagnostic pointing at this
  // compiler-generated op rather than the user-written BD.
  //
  // TODO: migrate this check to DMABDOp::verify() once BD iteration becomes
  // an explicit `aie.dma_bd` field and that early-return can be removed.
  if (getD0ZeroBefore() > 0x3F || getD0ZeroAfter() > 0x3F)
    return emitOpError("D0 pad_before/pad_after exceeds the [0:63] range.");
  if (getD1ZeroBefore() > 0x1F || getD1ZeroAfter() > 0x1F)
    return emitOpError("D1 pad_before/pad_after exceeds the [0:31] range.");
  if (getD2ZeroBefore() > 0xF || getD2ZeroAfter() > 0xF)
    return emitOpError("D2 pad_before/pad_after exceeds the [0:15] range.");
 
  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());
  }
  if (!targetModel.isShimNOCTile(getColumn(), getRow()) && getAxcache())
    return emitOpError("Only ShimTiles support AxCACHE configuration.");
 
  return success();
}
 
uint32_t AIEX::NpuWriteBdOp::getAxcacheOrDefault() {
  return getAxcache().value_or(AIE::getTargetModel(*this).getDefaultAxCache());
}
 
std::optional<uint32_t> AIEX::getConstantIntOperand(mlir::Value v) {
  if (std::optional<uint64_t> cst = getConstantInt64Operand(v))
    return static_cast<uint32_t>(*cst);
  return std::nullopt;
}
 
// Widthwise counterpart of getConstantIntOperand; see createConstantArgPlus.
std::optional<uint64_t> AIEX::getConstantInt64Operand(mlir::Value v) {
  // An arith.constant is read in place: the generic match folds the op, which
  // dominates translating a sequence of a million writes.
  if (auto constant = v.getDefiningOp<arith::ConstantOp>())
    if (auto attr = dyn_cast<IntegerAttr>(constant.getValue()))
      return attr.getValue().getZExtValue();
  mlir::APInt cst;
  if (!mlir::matchPattern(v, mlir::m_ConstantInt(&cst)))
    return std::nullopt;
  return cst.getZExtValue();
}
 
mlir::Value AIEX::createConstantI32(mlir::OpBuilder &builder,
                                    mlir::Location loc, uint32_t value) {
  return arith::ConstantOp::create(
      builder, loc, builder.getI32IntegerAttr(static_cast<int32_t>(value)));
}
 
mlir::Value AIEX::createConstantArgPlus(mlir::OpBuilder &builder,
                                        mlir::Location loc, uint64_t value) {
  if (value <= std::numeric_limits<uint32_t>::max())
    return createConstantI32(builder, loc, static_cast<uint32_t>(value));
  return arith::ConstantOp::create(
      builder, loc,
      IntegerAttr::get(builder.getIntegerType(64),
                       static_cast<int64_t>(value)));
}
 
//===----------------------------------------------------------------------===//
// Buffer-relative runtime writes
//===----------------------------------------------------------------------===//
 
// DataLayout's default queries abort for types without a builtin layout or a
// type interface; a legal memref element (such as another memref) need not have
// either. Check component types before querying composite builtin layouts too.
static bool hasQueryableStorageLayout(Type type) {
  if (type.isIntOrFloat() || isa<IndexType, DataLayoutTypeInterface>(type))
    return true;
  if (auto vector = dyn_cast<VectorType>(type))
    return hasQueryableStorageLayout(vector.getElementType());
  if (auto complex = dyn_cast<ComplexType>(type))
    return hasQueryableStorageLayout(complex.getElementType());
  return false;
}
 
// Match BufferOp's allocation convention without an unchecked shape product or
// a fixed-size conversion of scalable element storage.
static std::optional<uint64_t>
getStaticMemRefStorageBytes(MemRefType type, const DataLayout &layout) {
  if (!type || !type.hasStaticShape() ||
      !hasQueryableStorageLayout(type.getElementType()))
    return std::nullopt;
  if (auto vector = dyn_cast<VectorType>(type.getElementType())) {
    // The default vector layout uses a signed getNumElements(), then unsigned
    // innermost-dimension padding and bit products. Guard that query itself.
    if (vector.isScalable() || vector.getRank() == 0)
      return std::nullopt;
    int64_t elements = 1;
    for (int64_t dim : vector.getShape()) {
      auto product = llvm::checkedMul(elements, dim);
      if (!product)
        return std::nullopt;
      elements = *product;
    }
    uint64_t inner = vector.getShape().back();
    auto padded = llvm::checkedMulUnsigned(
        static_cast<uint64_t>(elements) / inner, llvm::PowerOf2Ceil(inner));
    llvm::TypeSize bits = layout.getTypeSizeInBits(vector.getElementType());
    if (!padded || bits.isScalable() ||
        !llvm::checkedMulUnsigned(*padded, bits.getFixedValue()))
      return std::nullopt;
  }
  llvm::TypeSize elementBytes = layout.getTypeSize(type.getElementType());
  if (elementBytes.isScalable())
    return std::nullopt;
  if (llvm::is_contained(type.getShape(), int64_t{0}))
    return 0;
  uint64_t bytes = elementBytes.getFixedValue();
  for (int64_t dim : type.getShape()) {
    auto product = llvm::checkedMulUnsigned(bytes, static_cast<uint64_t>(dim));
    if (!product)
      return std::nullopt;
    bytes = *product;
  }
  return bytes;
}
 
template <typename OpTy>
static LogicalResult verifyBufferWriteScope(OpTy op) {
  if (op.getBuffer() && !op->template getParentOfType<AIE::DeviceOp>())
    return op.emitOpError("must be inside a device when buffer is present");
  return success();
}
 
template <typename OpTy>
static LogicalResult
verifyBufferWrite(OpTy op, std::optional<uint32_t> wordOffset,
                  uint64_t writeBytes, AIE::NamedOpTableCollection &names) {
  auto bufferAttr = op.getBufferAttr();
  if (!bufferAttr)
    return success();
  if (!wordOffset)
    return op.emitOpError("address must be constant when buffer is present");
  auto device = op->template getParentOfType<AIE::DeviceOp>();
  if (!device)
    return op.emitOpError("must be inside a device when buffer is present");
  // The device's region verifier owns the index and has verified all child
  // invariants before dispatching these graph-dependent checks.
  auto buffer = names.lookupIn<AIE::BufferOp>(device, bufferAttr.getAttr());
  if (!buffer)
    return op.emitOpError()
           << "buffer '" << bufferAttr.getValue() << "' not found in device";
  auto bufferBytes = getStaticMemRefStorageBytes(buffer.getType(),
                                                 DataLayout::closest(buffer));
  if (!bufferBytes)
    return op.emitOpError(
        "referenced buffer must have a static, fixed-size storage extent "
        "representable in 64 bits");
 
  uint64_t offsetBytes = uint64_t{*wordOffset} * sizeof(uint32_t);
  auto end = llvm::checkedAddUnsigned(offsetBytes, writeBytes);
  if (!end)
    return op.emitOpError("buffer-relative write range overflows 64-bit bytes");
  if (*end > *bufferBytes)
    return op.emitOpError() << "write range [" << offsetBytes << ", " << *end
                            << ") exceeds buffer @" << bufferAttr.getValue()
                            << " size " << *bufferBytes << " bytes";
 
  auto tile = buffer.getTileLike();
  const auto &targetModel = device.getTargetModel();
  uint64_t memoryBytes = tile.isMemTile() ? targetModel.getMemTileSize()
                                          : targetModel.getLocalMemorySize();
  uint64_t start = offsetBytes;
  if (auto base = buffer.getAddress()) {
    if (*base < 0)
      return op.emitOpError("referenced buffer address must be nonnegative");
    start += static_cast<uint64_t>(*base);
    if (start % sizeof(uint32_t) != 0)
      return op.emitOpError() << "resolved buffer write address " << start
                              << " is not 4-byte aligned";
  }
  // An unassigned base can only move the interval upward. Check this necessary
  // bound now, and the resolved interval when address assignment runs.
  auto memoryEnd = llvm::checkedAddUnsigned(start, writeBytes);
  if (!memoryEnd)
    return op.emitOpError("resolved buffer write range overflows 64-bit bytes");
  if (*memoryEnd > memoryBytes)
    return op.emitOpError()
           << "write range [" << start << ", " << *memoryEnd
           << ") exceeds tile data memory size " << memoryBytes << " bytes";
  return success();
}
 
//===----------------------------------------------------------------------===//
// NpuWrite32Op
//===----------------------------------------------------------------------===//
 
LogicalResult AIEX::NpuWrite32Op::verify() {
  return verifyBufferWriteScope(*this);
}
 
LogicalResult
AIEX::NpuWrite32Op::verifyNamedOpUses(AIE::NamedOpTableCollection &names) {
  if (!getBuffer())
    return success();
  return verifyBufferWrite(*this, getConstantIntOperand(getAddress()),
                           sizeof(uint32_t), names);
}
 
template <typename T>
static std::optional<uint32_t>
getAbsoluteAddress(T *op, uint32_t addressOffset,
                   const AIE::NamedOpTable *names) {
  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 (auto bufferSym = op->getBuffer()) {
    AIE::BufferOp buffer =
        names ? names->lookup<AIE::BufferOp>(op->getBufferAttr().getAttr())
              : AIE::lookupNamedOpIn<AIE::BufferOp>(device, *bufferSym);
    if (!buffer) {
      op->emitError() << "buffer '" << *bufferSym << "' not found in device";
      return std::nullopt;
    }
 
    auto bufferAddress = buffer.getAddress();
    if (!bufferAddress) {
      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>(*bufferAddress) +
              addressOffset * sizeof(uint32_t);
    address = ((col & 0xff) << tm.getColumnShift()) |
              ((row & 0xff) << tm.getRowShift()) | (address & 0xfffff);
  } else { // otherwise, the given address is absolute
    address = addressOffset;
    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(const AIE::NamedOpTable *names) {
  std::optional<uint32_t> addressOffset = getConstantIntOperand(getAddress());
  if (!addressOffset)
    return std::nullopt;
  return ::getAbsoluteAddress(this, *addressOffset, names);
}
 
//===----------------------------------------------------------------------===//
// NpuAddressPatchOp
//===----------------------------------------------------------------------===//
 
LogicalResult AIEX::NpuAddressPatchOp::verify() {
  // A runtime register address (addr_val) is meaningful only on the EmitC (C++
  // TXN) target; the constant `addr` still carries the fallback / static value.
  // The static binary target checks for the operand and diagnoses it there.
  if (getAddrVal() && !getAddrVal().getType().isInteger(32))
    return emitOpError("addr_val must be an i32 value");
  if (getArgIdx().has_value() == static_cast<bool>(getBuffer()))
    return emitOpError("must name the host buffer either by the 'arg_idx' "
                       "attribute or by the 'buffer' operand, not both and not "
                       "neither");
  return success();
}
 
//===----------------------------------------------------------------------===//
// NpuUpdateFromScratchpadOp
//===----------------------------------------------------------------------===//
 
std::optional<uint32_t> AIEX::NpuUpdateFromScratchpadOp::getAbsoluteAddress(
    const AIE::NamedOpTable *names) {
  return ::getAbsoluteAddress(this, getAddress(), names);
}
 
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) << ").";
 
  // NpuCreateScratchpadOp::verify() bounds the state_table_idx.
  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.";
  }
 
  // Scratchpads are far rarer than the updates they bound, and an op verifier
  // runs after every pass, so the block scan belongs on this side.
  uint32_t numEntries = getSize() / 4;
  if (Block *block = (*this)->getBlock()) {
    for (auto updateOp : block->getOps<AIEX::NpuUpdateFromScratchpadOp>()) {
      if (updateOp.getStateTableIdx() < numEntries)
        continue;
      return updateOp.emitOpError("state_table_idx ")
             << static_cast<uint32_t>(updateOp.getStateTableIdx())
             << " is out of bounds for scratchpad of size " << getSize()
             << " bytes (" << numEntries << " entries) created by "
             << (*this)->getName() << ".";
    }
  }
 
  // At most one create_scratchpad may appear per runtime sequence. Walk the
  // parent RuntimeSequenceOp to check; only report from the duplicate (i.e.
  // the op that is NOT the first occurrence) to avoid emitting the same error
  // twice.
  auto runtimeSeq = getOperation()->getParentOfType<AIE::RuntimeSequenceOp>();
  if (!runtimeSeq) {
    return success();
  }
 
  NpuCreateScratchpadOp firstSeen;
  runtimeSeq.walk([&](NpuCreateScratchpadOp op) {
    firstSeen = op;
    return WalkResult::interrupt();
  });
  if (firstSeen != *this) {
    InFlightDiagnostic diag =
        emitOpError("only one 'aiex.npu.create_scratchpad' is allowed per "
                    "runtime sequence");
    diag.attachNote(firstSeen.getLoc())
        << "previous 'aiex.npu.create_scratchpad' here";
    return diag;
  }
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// NpuMaskWrite32Op
//===----------------------------------------------------------------------===//
 
LogicalResult AIEX::NpuMaskWrite32Op::verify() {
  return verifyBufferWriteScope(*this);
}
 
LogicalResult
AIEX::NpuMaskWrite32Op::verifyNamedOpUses(AIE::NamedOpTableCollection &names) {
  if (!getBuffer())
    return success();
  return verifyBufferWrite(*this, getConstantIntOperand(getAddress()),
                           sizeof(uint32_t), names);
}
 
std::optional<uint32_t>
AIEX::NpuMaskWrite32Op::getAbsoluteAddress(const AIE::NamedOpTable *names) {
  std::optional<uint32_t> addressOffset = getConstantIntOperand(getAddress());
  if (!addressOffset)
    return std::nullopt;
  return ::getAbsoluteAddress(this, *addressOffset, names);
}
 
//===----------------------------------------------------------------------===//
// NpuMaskPollOp
//===----------------------------------------------------------------------===//
 
LogicalResult AIEX::NpuMaskPollOp::verify() {
  return verifyBufferWriteScope(*this);
}
 
LogicalResult
AIEX::NpuMaskPollOp::verifyNamedOpUses(AIE::NamedOpTableCollection &names) {
  if (!getBuffer())
    return success();
  return verifyBufferWrite(*this, getConstantIntOperand(getAddress()),
                           sizeof(uint32_t), names);
}
 
std::optional<uint32_t>
AIEX::NpuMaskPollOp::getAbsoluteAddress(const AIE::NamedOpTable *names) {
  std::optional<uint32_t> addressOffset = getConstantIntOperand(getAddress());
  if (!addressOffset)
    return std::nullopt;
  return ::getAbsoluteAddress(this, *addressOffset, names);
}
 
//===----------------------------------------------------------------------===//
// NpuBlockWriteOp
//===----------------------------------------------------------------------===//
 
LogicalResult AIEX::NpuBlockWriteOp::verify() {
  return verifyBufferWriteScope(*this);
}
 
LogicalResult
AIEX::NpuBlockWriteOp::verifyNamedOpUses(AIE::NamedOpTableCollection &names) {
  if (!getBuffer())
    return success();
  auto type = dyn_cast<MemRefType>(getData().getType());
  DataLayout layout = DataLayout::closest(*this);
  auto bytes = getStaticMemRefStorageBytes(type, layout);
  if (!bytes)
    return emitOpError("payload must have a static, fixed-size storage extent "
                       "representable in 64 bits when buffer is present");
  llvm::TypeSize elementBits = layout.getTypeSizeInBits(type.getElementType());
  if (elementBits.isScalable() || elementBits.getFixedValue() != 32)
    return emitOpError(
        "payload must have 32-bit elements when buffer is present");
  return verifyBufferWrite(*this, getAddress(), *bytes, names);
}
 
std::optional<uint32_t>
AIEX::NpuBlockWriteOp::getAbsoluteAddress(const AIE::NamedOpTable *names) {
  return ::getAbsoluteAddress(this, getAddress(), names);
}
 
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
//===----------------------------------------------------------------------===//
 
AIE::DMABDOp AIEX::DMAConfigureTaskOp::getFirstBd() {
  Region &body = getBody();
  if (body.empty()) {
    return nullptr;
  }
  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 nullptr;
  }
  return *bd_ops.begin();
}
 
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 &block : body) {
    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();
}
 
// Enforce the per-BD ND access-pattern limit for BDs nested inside a
// runtime-sequence DMA task. The AIE::DMABDOp verifier skips these BDs (their
// parent is a DMA task op, not a *DMAOp), so this is the only check of the BD
// dimension count on the runtime-sequence path.
//
// Every AIE2/AIE2P DMA BD register file carries getBDMaxDims ND address
// dimensions (D0..) plus one separate iteration/repeat dimension: a core/shim
// BD has D0..D2 + iteration, a MemTile BD has D0..D3 + iteration. On this path
// aiex.shim_dma_single_bd_task hoists the leading tap dimension into that
// iteration register, so a shim/core BD may carry one dimension beyond its ND
// access limit (3 + 1). A MemTile is not given the +1: AIEDMATasksToNPU maps
// the 4th task dimension onto the iteration register for every tile type and
// caps the total at 4, which the MemTile's 4 ND dimensions already reach. Both
// branches therefore land on the same uniform 4-dimension cap enforced later by
// AIEDMATasksToNPU.
//
// A BD whose sizes and strides are all constant may give more: every dimension
// past the third is an iteration dimension, and aie-decompose-large-dma-bd
// splits them off into descriptors of 4. With a runtime value among them it
// cannot, so the cap applies here.
//
// A BD's `iteration` attribute claims that iteration register itself, so its
// sizes/strides must leave it free: one dimension fewer, with no splitting.
static LogicalResult
verifyTaskBDDimensions(const AIE::AIETargetModel &targetModel, int col, int row,
                       Region &body) {
  size_t maxNDims = targetModel.getBDMaxDims(col, row);
  if (!targetModel.isMemTile(col, row))
    ++maxNDims; // leading dim is hoisted into the iteration/repeat register
  LogicalResult result = success();
  body.walk([&](AIE::DMABDOp bd) {
    // The BD's own verifier skips it here, so nothing has yet established that
    // its mixed sizes/strides lists are safe to read.
    if (failed(bd.verifyMixedSizesAndStrides())) {
      result = failure();
      return;
    }
    if (failed(
            AIE::verifyDMABDIteration(bd, targetModel.getTileType(col, row)))) {
      result = failure();
      return;
    }
    size_t numDims = bd.getMixedSizes().size();
    if (bd.getIteration()) {
      if (numDims + 1 > maxNDims) {
        bd.emitOpError() << "Cannot give more than "
                         << std::to_string(maxNDims - 1)
                         << " dimensions for step sizes and wraps alongside "
                            "the iteration attribute on this tile (got "
                         << std::to_string(numDims) << " dimensions).";
        result = failure();
      }
      return;
    }
    auto isConstant = [](OpFoldResult v) {
      return getConstantIntValue(v).has_value();
    };
    if (numDims > maxNDims && llvm::all_of(bd.getMixedSizes(), isConstant) &&
        llvm::all_of(bd.getMixedStrides(), isConstant))
      return;
    if (numDims > maxNDims) {
      bd.emitOpError() << "Cannot give more than " << std::to_string(maxNDims)
                       << " dimensions for step sizes and wraps on this tile "
                          "(got "
                       << std::to_string(numDims) << " dimensions).";
      result = failure();
    }
  });
  return result;
}
 
LogicalResult AIEX::DMAConfigureTaskOp::verify() {
  const AIE::AIETargetModel &targetModel = AIE::getTargetModel(getOperation());
  // Skip the per-BD dimension check on an unplaced (logical) tile: the ND limit
  // is a function of the tile's placed coordinates, which are not yet known.
  // The verifier runs again on the concrete tile once placement resolves it.
  std::optional<int> col = getTileLike().tryGetCol();
  std::optional<int> row = getTileLike().tryGetRow();
  if (col && row &&
      failed(verifyTaskBDDimensions(targetModel, *col, *row, getBody())))
    return failure();
  Region &body = getBody();
  // 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
  if (failed(AIE::verifyAssignedPacket(*this, getPacketAttr())))
    return failure();
  LogicalResult result = success();
  bool taskHasPacket = getPacket().has_value();
  llvm::SmallVector<AIE::DMABDOp> bds;
  for (auto &block : body) {
    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());
 
    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();
      }
      // DMABDOp::verify skips task BDs, so validate their packet header and
      // out_of_order_id here too.
      if (failed(AIE::verifyAssignedPacket(
              bd, bd.getPacket().value_or(AIE::PacketInfoAttr()))) ||
          failed(AIE::verifyDMABDOutOfOrderId(bd, taskHasPacket)) ||
          failed(bd.verifyTransferInBounds()))
        result = failure();
      bds.push_back(bd);
    });
  }
  if (getOutOfOrder()) {
    // Out-of-order mode rejects task completion token (bits are aliased).
    if (getIssueToken()) {
      auto err =
          emitOpError("out_of_order channel cannot issue a completion token");
      err.attachNote() << "set issue_token = false on this out-of-order task";
      result = failure();
    }
    if (failed(AIE::verifyOutOfOrderChannel(
            getOperation(), getDirection(), getOutOfOrder(),
            llvm::ArrayRef<AIE::DMABDOp>(bds),
            /*packetEnabledByContext=*/taskHasPacket)))
      result = failure();
  }
  return result;
}
 
LogicalResult AIEX::DMAStartTaskOp::verify() {
  if (getRepeatCountAttr() && getRepeatCountVal())
    return emitOpError(
        "takes repeat_count or a runtime repeat count, not both");
  if (IntegerAttr rc = getRepeatCountAttr(); rc && rc.getInt() < 0)
    return emitOpError("repeat_count must be non-negative, got ")
           << rc.getInt();
  if (hasRepeatParameter() && (getRepeatCountAttr() || getRepeatCountVal()))
    return emitOpError("takes a repeat_parameter or a repeat count, not both");
  if (getRepeatParameterAttr() && getRepeatStateTableIdxAttr())
    return emitOpError("takes repeat_parameter or repeat_state_table_idx, "
                       "not both");
  return success();
}
 
// A BD's static length in bytes: the explicit or implicit (whole-buffer)
// length, scaled by the element width. Nullopt for a runtime length, or one a
// scratchpad parameter lengthens at run time (AIEOps.td's "Runtime length").
static std::optional<int64_t> staticBdLenBytes(AIE::DMABDOp bd) {
  if (bd.getLengthParameterAttr() || bd.getLengthStateTableIdxAttr())
    return std::nullopt;
  int64_t elems;
  if (std::optional<int32_t> len = bd.getConstantLen()) {
    // Zero-extended, as DMABDOp::getLenInBytes() reads the length.
    elems = static_cast<uint32_t>(*len);
  } else {
    // With runtime sizes/strides, an omitted len is their inner extent
    // (AIEDMATasksToNPU), not the whole buffer.
    if (bd.hasLen() || !bd.getSizes().empty() || !bd.getStrides().empty())
      return std::nullopt;
    auto buffer = dyn_cast<MemRefType>(bd.getBuffer().getType());
    if (!buffer || !buffer.hasStaticShape())
      return std::nullopt;
    elems = buffer.getNumElements();
  }
  return llvm::checkedMul(
      elems, static_cast<int64_t>(bd.getBufferElementTypeWidthInBytes()));
}
 
// Bytes a task streams per dispatch, or nullopt when any extent, repeat or
// start count is not static, or the total overflows.
static std::optional<int64_t>
staticBytesStreamed(AIEX::DMAConfigureTaskForOp task) {
  int64_t starts = 0;
  for (Operation *user : task.getResult().getUsers()) {
    auto start = dyn_cast<AIEX::DMAStartTaskOp>(user);
    if (!start) {
      if (isa<AIEX::DMAAwaitTaskOp, AIEX::DMAFreeTaskOp>(user))
        continue;
      // The task value is forwarded (a branch operand, a call, ...), so its
      // starts elsewhere can't be counted.
      return std::nullopt;
    }
    // A start outside the task's own block may run under a loop or a
    // mutually-exclusive branch, either of which this static count can't see.
    if (start->getBlock() != task->getBlock() || start.getRepeatCountAttr() ||
        start.getRepeatCountVal() || start.hasRepeatParameter())
      return std::nullopt;
    ++starts;
  }
  if (starts == 0)
    return 0;
  std::optional<int64_t> repeat =
      getConstantIntValue(task.getRepeatCountValue());
  if (!repeat || *repeat < 0)
    return std::nullopt;
 
  Region &body = task.getBody();
  if (body.empty())
    return std::nullopt;
  // The queue starts the chain at the entry block, so only blocks reachable
  // from it stream. A cycle runs indefinitely, and a next_bd_id continues the
  // chain outside the body.
  std::optional<int64_t> moved = 0;
  for (auto scc = llvm::scc_begin(&body.front()); !scc.isAtEnd(); ++scc) {
    if (scc.hasCycle())
      return std::nullopt;
    for (AIE::DMABDOp bd : scc->front()->getOps<AIE::DMABDOp>()) {
      std::optional<int64_t> len = staticBdLenBytes(bd);
      if (!len || bd.getIteration() || bd.getNextBdId())
        return std::nullopt;
      moved = llvm::checkedAdd(*moved, *len);
      if (!moved)
        return std::nullopt;
    }
  }
  std::optional<int64_t> runs = llvm::checkedAdd(*repeat, int64_t{1});
  if (runs)
    runs = llvm::checkedMul(*runs, starts);
  return runs ? llvm::checkedMul(*moved, *runs) : std::nullopt;
}
 
namespace {
// A far end's cyclic chain of BDs over one object: each segment's BD, in
// order, sent `copies` times in a row.
struct BdChain {
  SmallVector<int64_t> segmentBytes;
  int64_t copies;
  int64_t objectBytes;
 
  // Nullopt when the object's bytes overflow or a segment is empty.
  static std::optional<BdChain> get(ArrayRef<int64_t> segmentBytes,
                                    int64_t copies) {
    if (segmentBytes.size() == 1)
      copies = 1;
    std::optional<int64_t> objectBytes = 0;
    for (int64_t bytes : segmentBytes) {
      std::optional<int64_t> run = llvm::checkedMul(bytes, copies);
      if (bytes <= 0 || copies <= 0 || !run)
        return std::nullopt;
      objectBytes = llvm::checkedAdd(*objectBytes, *run);
      if (!objectBytes)
        return std::nullopt;
    }
    if (segmentBytes.empty())
      return std::nullopt;
    return BdChain{SmallVector<int64_t>(segmentBytes), copies, *objectBytes};
  }
 
  // True when `bytes` through the chain stop where one of its BDs ends.
  bool endsOnBd(int64_t bytes) const {
    int64_t rest = bytes % objectBytes;
    int64_t start = 0;
    for (int64_t segment : segmentBytes) {
      if (rest < start + segment * copies)
        return (rest - start) % segment == 0;
      start += segment * copies;
    }
    return true;
  }
 
  bool operator==(const BdChain &other) const {
    return segmentBytes == other.segmentBytes && copies == other.copies;
  }
};
} // namespace
 
// The far end of a shim transfer on an objectFIFO receives (or sends) whole
// objects through BDs sized by the object type, or by its link pool's
// segments, so transfers ending mid-BD never complete it and the fifo hangs.
// Several transfers may share an object (concatenating host buffers into it, or
// splitting it across them), so it is the runtime sequence's static total
// through the fifo that must stop at a BD boundary.
static LogicalResult
verifyTaskCoversWholeObjects(AIEX::DMAConfigureTaskForOp task,
                             AIE::ObjectFifoCreateOp fifo) {
  if (fifo.getPadDimensions() || fifo.usesStream())
    return success();
  std::optional<int64_t> moved = staticBytesStreamed(task);
  if (!moved || *moved == 0)
    return success();
  DataLayout layout = DataLayout::closest(task);
  // Bytes an object holds, or nullopt when that overflows.
  auto objectBytesOf = [&](Type fifoType) -> std::optional<int64_t> {
    MemRefType elem = cast<AIE::AIEObjectFifoType>(fifoType).getElementType();
    return llvm::checkedMul(
        elem.getNumElements(),
        static_cast<int64_t>(layout.getTypeSize(elem.getElementType())));
  };
  // A linked fifo's far end on the link's shared tile is a chain of BDs over
  // the link's pool, one per pool segment, sized by the pool owner whatever
  // this fifo's own (or consumer) type. A join's or distribute's combined end
  // steps through every segment in order, each small end through its own, and
  // a plain link's single segment is the whole pool object. The link's
  // repeat_count sends each of its output's segment BDs that many times in a
  // row. Every link the fifo meets has an end the transfer must complete,
  // one per tile when links on several tiles read a broadcast.
  SmallVector<Value> linkTiles;
  SmallVector<BdChain> chains;
  if (AIE::DeviceOp device = fifo->getParentOfType<AIE::DeviceOp>())
    for (AIE::ObjectFifoLinkOp link : device.getOps<AIE::ObjectFifoLinkOp>()) {
      std::vector<AIE::ObjectFifoCreateOp> ins = link.getInputObjectFifos();
      std::vector<AIE::ObjectFifoCreateOp> outs = link.getOutputObjectFifos();
      if (!llvm::is_contained(ins, fifo) && !llvm::is_contained(outs, fifo))
        continue;
      std::optional<Value> sharedTile = link.getOptionalSharedTile();
      if (ins.empty() || outs.empty() || !sharedTile)
        return success();
      linkTiles.push_back(*sharedTile);
      AIE::ObjectFifoCreateOp owner = link.getPoolOwner();
      MemRefType pool =
          cast<AIE::AIEObjectFifoType>(owner.getElemType()).getElementType();
      auto width =
          static_cast<int64_t>(layout.getTypeSize(pool.getElementType()));
      SmallVector<std::pair<int64_t, int64_t>> segments =
          link.getPoolSegments();
      if (fifo != owner && (link.isJoin() || link.isDistribute())) {
        std::vector<AIE::ObjectFifoCreateOp> &many = link.isJoin() ? ins : outs;
        segments = {segments[llvm::find(many, fifo) - many.begin()]};
      }
      SmallVector<int64_t> segmentBytes;
      for (auto [offset, size] : segments) {
        std::optional<int64_t> bytes = llvm::checkedMul(size, width);
        if (!bytes)
          break;
        segmentBytes.push_back(*bytes);
      }
      int64_t copies = llvm::is_contained(outs, fifo)
                           ? link.getRepeatCount().value_or(1)
                           : 1;
      std::optional<BdChain> chain;
      if (segmentBytes.size() == segments.size())
        chain = BdChain::get(segmentBytes, copies);
      if (chain && !llvm::is_contained(chains, *chain))
        chains.push_back(*chain);
    }
  // Every far-end BD chain this transfer must leave at a BD boundary: the
  // producer's, for a drain, and each consumer's other than the shim's own. A
  // broadcast stalls on its slowest consumer, so a drain bounds what the
  // others receive too. An end off every link tile is one BD over its own
  // object.
  auto addChain = [&](Value tile, Type fifoType) {
    if (llvm::is_contained(linkTiles, tile))
      return;
    std::optional<BdChain> chain;
    if (std::optional<int64_t> bytes = objectBytesOf(fifoType))
      chain = BdChain::get({*bytes}, 1);
    if (chain && !llvm::is_contained(chains, *chain))
      chains.push_back(*chain);
  };
  auto producer =
      dyn_cast_or_null<AIE::TileLike>(fifo.getProducerTile().getDefiningOp());
  if (!producer)
    return success();
  if (!producer.isShimTile())
    addChain(fifo.getProducerTile(), fifo.getElemType());
  for (Value consumer : fifo.getConsumerTiles()) {
    auto tile = dyn_cast_or_null<AIE::TileLike>(consumer.getDefiningOp());
    if (!tile || !tile.isShimTile())
      addChain(consumer, fifo.getConsumerElemTypeOrDefault());
  }
  if (chains.empty())
    return success();
 
  // Transfers in a loop or branch have no static count to sum.
  auto sequence = task->getParentOfType<AIE::RuntimeSequenceOp>();
  if (!sequence || task->getBlock() != &sequence.getBody().front())
    return success();
  // A dma_channel_reset_for re-arms the fifo's far ends and drops a partly
  // filled object. It doesn't block, so which bytes land before it isn't
  // static, and the transfers around it can't be summed.
  if (llvm::any_of(task->getBlock()->getOps<AIEX::DmaChannelResetForOp>(),
                   [&](AIEX::DmaChannelResetForOp reset) {
                     return reset.getObjfifoAttr() == task.getAlloc();
                   }))
    return success();
  // Segment boundaries don't add up the way whole objects do, so the
  // sequence's total is checked whatever this transfer moves, once: by the
  // first transfer through the fifo.
  for (Operation *prev = task->getPrevNode(); prev; prev = prev->getPrevNode())
    if (auto other = dyn_cast<AIEX::DMAConfigureTaskForOp>(prev);
        other && other.getAlloc() == task.getAlloc() &&
        staticBytesStreamed(other) != 0)
      return success();
  std::optional<int64_t> total = 0;
  int64_t sharing = 0;
  WalkResult walk = sequence.walk([&](Operation *op) {
    // A dma_memcpy_nd or BD chain through the fifo also shares its objects;
    // its extent is not summed here, so the sequence's total is unknown.
    if (auto memcpy = dyn_cast<AIEX::NpuDmaMemcpyNdOp>(op))
      return memcpy.getMetadata() == task.getAlloc() ? WalkResult::interrupt()
                                                     : WalkResult::advance();
    if (auto chain = dyn_cast<AIEX::DMAStartBdChainForOp>(op))
      return chain.getAllocAttr() == task.getAlloc() ? WalkResult::interrupt()
                                                     : WalkResult::advance();
    auto other = dyn_cast<AIEX::DMAConfigureTaskForOp>(op);
    if (!other || other.getAlloc() != task.getAlloc())
      return WalkResult::advance();
    std::optional<int64_t> bytes = staticBytesStreamed(other);
    // A task that is never started streams nothing, wherever it is.
    if (bytes == 0)
      return WalkResult::advance();
    if (!bytes || other->getBlock() != task->getBlock())
      return WalkResult::interrupt();
    total = llvm::checkedAdd(*total, *bytes);
    if (!total)
      return WalkResult::interrupt();
    ++sharing;
    return WalkResult::advance();
  });
  if (walk.wasInterrupted())
    return success();
  const BdChain *partial = llvm::find_if(
      chains, [&](const BdChain &chain) { return !chain.endsOnBd(*total); });
  if (partial == chains.end())
    return success();
 
  InFlightDiagnostic error =
      task.emitOpError() << "moves " << *moved << " bytes through objectFIFO @"
                         << fifo.getSymName() << ", which";
  if (sharing > 1)
    error << ", with this runtime sequence's other transfers through it ("
          << *total << " bytes in all),";
  error << " is not a whole number of its " << partial->objectBytes
        << "-byte objects";
  if (partial->segmentBytes.size() == 1)
    return error << "; the far end's buffer descriptor is sized by the "
                    "object, so a partial one never completes";
  std::string segments;
  llvm::raw_string_ostream os(segments);
  llvm::interleaveComma(partial->segmentBytes, os);
  error << " and doesn't end on a segment boundary (segments of " << segments
        << " bytes";
  if (partial->copies > 1)
    error << ", each sent " << partial->copies << " times in a row";
  return error << "); the far end's buffer descriptors are sized by the "
                  "segments, so a partial one never completes";
}
 
// Resolving the allocation symbol through the collection keeps the lookup off
// the device's linear symbol scan, which a per-op verifier repeats after every
// pass.
LogicalResult AIEX::DMAConfigureTaskForOp::verifySymbolUses(
    SymbolTableCollection &symbolTable) {
  AIE::DeviceOp dev = getOperation()->getParentOfType<AIE::DeviceOp>();
  if (!dev)
    return success();
  Operation *target =
      symbolTable.lookupSymbolIn(dev, getAlloc().getRootReference());
  if (!target)
    return emitOpError() << getAlloc() << " does not name a symbol";
  // The tile is known once the name is a shim DMA allocation or a route
  // endpoint, so the per-BD dimension limit can be enforced before the task
  // becomes a concrete dma_configure_task. Before the objectFIFO lowering
  // the name is the fifo's own, and the check waits.
  Value tileValue;
  LogicalResult named =
      TypeSwitch<Operation *, LogicalResult>(target)
          .Case([&](AIE::ShimDMAAllocationOp alloc) {
            tileValue = alloc.getTile();
            return success();
          })
          .Case([&](AIE::RouteEndpointOp endpoint) -> LogicalResult {
            tileValue = endpoint.getTile();
            if (endpoint.getBundle() == AIE::WireBundle::DMA)
              return success();
            return emitOpError() << getAlloc() << " names a "
                                 << stringifyWireBundle(endpoint.getBundle())
                                 << " port, not a DMA channel";
          })
          .Case([&](AIE::ObjectFifoCreateOp fifo) {
            return verifyTaskCoversWholeObjects(*this, fifo);
          })
          .Default([&](Operation *) -> LogicalResult {
            return emitOpError() << getAlloc()
                                 << " must name an aie.shim_dma_allocation, an "
                                    "aie.route_endpoint or an aie.objectfifo";
          });
  if (failed(named))
    return failure();
  // An allocation or endpoint on an unplaced (logical) tile waits for
  // placement to substitute a real one.
  auto tile = tileValue ? tileValue.getDefiningOp<AIE::TileOp>() : nullptr;
  if (!tile)
    return success();
  const AIE::AIETargetModel &targetModel = AIE::getTargetModel(getOperation());
  return verifyTaskBDDimensions(targetModel, tile.getCol(), tile.getRow(),
                                getBody());
}
 
//===----------------------------------------------------------------------===//
// 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 (std::optional<int64_t> value = getConstantIntValue(getValue())) {
    if (*value < 0)
      return emitOpError("Lock value must be non-negative");
    if (*value > 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();
  }
  // A logical tile with a coordinate still unpinned waits for the
  // post-placement re-verify for its range and address checks, as
  // UsesAreAccessible's do.
  std::optional<int> col = lockOp.getTileLike().tryGetCol();
  std::optional<int> row = lockOp.getTileLike().tryGetRow();
  if (!col || !row)
    return success();
 
  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, AIE::TileID{*col, *row})) {
    return emitOpError("Invalid lock ID and tile combination when trying to "
                       "retrieve the local lock address.");
  }
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// DMABdPoolPopOp / DMABdPoolPushOp
//===----------------------------------------------------------------------===//
 
static LogicalResult verifyBdPoolPartition(Operation *op, int col, int row,
                                           uint32_t begin, uint32_t end) {
  uint32_t numBds = AIE::getTargetModel(op).getNumBDs(col, row);
  if (begin >= end || end > numBds)
    return op->emitOpError()
           << "partition [" << begin << ", " << end
           << ") is not a nonempty range of the " << numBds
           << " buffer descriptors on tile (" << col << ", " << row << ")";
  return success();
}
 
LogicalResult AIEX::DMABdPoolPopOp::verify() {
  return verifyBdPoolPartition(*this, getColumn(), getRow(), getBdBegin(),
                               getBdEnd());
}
 
LogicalResult AIEX::DMABdPoolPushOp::verify() {
  return verifyBdPoolPartition(*this, getColumn(), getRow(), getBdBegin(),
                               getBdEnd());
}
 
//===----------------------------------------------------------------------===//
// DmaChannelResetOp
//===----------------------------------------------------------------------===//
 
LogicalResult AIEX::DmaChannelResetOp::verify() {
  const auto &targetModel = AIE::getTargetModel(*this);
 
  auto tile = dyn_cast_or_null<AIE::TileOp>(getTile().getDefiningOp());
  if (!tile)
    return emitOpError() << "tile operand must be produced by an aie.tile op";
  int col = tile.getCol();
  int row = tile.getRow();
  // Only core and mem tiles have a per-channel DMA reset bit. The shim NOC DMA
  // control register has no reset field (aie-rt's Aie2PShimDmaChProp sets
  // Reset.Mask = 0); bit 1 there is PAUSE_MEM, not RESET. Rejecting shim keeps
  // the op honest about what it can lower, matching aie-rt's own
  // XAie_DmaChannelReset, which errors on SHIMNOC/SHIMPL tiles.
  if (!targetModel.isCoreTile(col, row) && !targetModel.isMemTile(col, row))
    return emitOpError() << "tile (" << col << ", " << row
                         << ") has no DMA channel reset (only core and mem "
                            "tiles do; shim NOC DMA has no reset bit)";
 
  // Number of DMA channels on this tile in this direction. Mirrors
  // TileOp::getNumSource/DestConnections(WireBundle::DMA): the switchbox
  // direction is reversed relative to the DMA direction.
  uint32_t numChannels = getDirection() == AIE::DMAChannelDir::S2MM
                             ? targetModel.getNumDestSwitchboxConnections(
                                   col, row, AIE::WireBundle::DMA)
                             : targetModel.getNumSourceSwitchboxConnections(
                                   col, row, AIE::WireBundle::DMA);
  if (getChannel() >= numChannels)
    return emitOpError() << "channel " << getChannel()
                         << " out of range for this tile and direction (tile "
                            "has "
                         << numChannels << " DMA channel(s) in this direction)";
 
  return success();
}
 
LogicalResult AIEX::DmaChannelResetForOp::verify() {
  // Deferring verifier, in the shape of DMAConfigureTaskForOp::verify: the
  // referenced symbol is only resolvable once the objectFIFO lowering has run.
  // It names an aie.objectfifo, then that fifo's shim endpoint while the
  // lowering is in flight, and finally its aie.objectfifo_rearm_binding. If
  // none is resolvable yet, defer -- a later pass will resolve it.
  AIE::DeviceOp dev = getOperation()->getParentOfType<AIE::DeviceOp>();
  if (!dev)
    return success();
  // The resident re-arm relies on the aie2p behavior that a DMA channel has no
  // enable bit, so the only way to restart it is a START_QUEUE push. AIE1 DMA
  // channels have an enable bit and are armed differently, so the trio this op
  // lowers to would not re-arm them correctly.
  if (AIE::getTargetModel(*this).getTargetArch() == AIE::AIEArch::AIE1)
    return emitOpError() << "is not supported on AIE1 devices (the resident "
                            "re-arm relies on the aie2p start-queue push)";
  Operation *target = dev.lookupSymbol(getObjfifo());
  if (!target)
    return success(); // symbol resolved during a later pass; defer the check
  if (!isa<AIE::ObjectFifoCreateOp, AIE::RouteEndpoint,
           AIE::ObjectFifoRearmBindingOp>(target))
    return emitOpError() << "'" << getObjfifo()
                         << "' must reference an aie.objectfifo (or its "
                            "aie.objectfifo_rearm_binding)";
  return success();
}
 
//===----------------------------------------------------------------------===//
// CoreResetOp
//===----------------------------------------------------------------------===//
 
LogicalResult AIEX::CoreResetOp::verify() {
  const auto &targetModel = AIE::getTargetModel(*this);
 
  // The op lowers to an NPU control-packet write; the runtime sequence has no
  // meaning on AIE1. Reject it explicitly, as SetLockOp does.
  if (targetModel.getTargetArch() == AIE::AIEArch::AIE1)
    return emitOpError("aiex.core_reset is not supported on AIE1.");
 
  auto tile = dyn_cast_or_null<AIE::TileOp>(getTile().getDefiningOp());
  if (!tile)
    return emitOpError() << "tile operand must be produced by an aie.tile op";
  int col = tile.getCol();
  int row = tile.getRow();
  // The tile coordinates are bounded by aie.tile's own verifier, so this op
  // does not re-check them for range.
 
  // Only core tiles have a CORE_CONTROL register with a reset bit. Mem and shim
  // tiles have no compute core, so there is nothing valid to lower to. This
  // matches aie-rt's XAie_CoreReset, which errors on any tile that is not an
  // AIE (core) tile.
  if (!targetModel.isCoreTile(col, row))
    return emitOpError() << "tile (" << col << ", " << row
                         << ") has no core to reset (only core tiles have a "
                            "CORE_CONTROL register)";
 
  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;
}
 
//===----------------------------------------------------------------------===//
// ReadScratchpadParameterOp
//===----------------------------------------------------------------------===//
 
LogicalResult AIEX::ReadScratchpadParameterOp::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 moduleOp = (*this)->getParentOfType<ModuleOp>();
  if (!moduleOp ||
      !moduleOp.lookupSymbol<AIEX::ScratchpadParameterOp>(getParameter())) {
    return emitOpError("references unknown parameter '")
           << getParameter()
           << "' (aiex.scratchpad_parameter ops are declared at module scope)";
  }
  if (getResult().getType().isF32()) {
    return emitOpError(
        "f32 parameters are not supported: the scratchpad encoding zeroes "
        "the top 2 bits, which clobbers the sign bit and top exponent bit "
        "of an f32. Use bf16 or an integer type up to i32 instead.");
  }
  return success();
}
 
LogicalResult AIEX::WaitScratchpadParametersOp::verify() {
  if (!(*this)->getParentOfType<AIE::CoreOp>()) {
    return emitOpError("must be inside an aie.core");
  }
  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(AIE::NamedOpTableCollection *tables) {
  AIEX::ConfigureOp configureOp =
      getOperation()->getParentOfType<AIEX::ConfigureOp>();
  if (!configureOp) {
    return nullptr;
  }
  AIE::DeviceOp referencedDevice = configureOp.getReferencedDeviceOp();
  if (!referencedDevice) {
    return nullptr;
  }
 
  StringAttr name = getRuntimeSequenceSymbolAttr().getAttr();
  if (!tables) {
    return AIE::lookupNamedOpIn<AIE::RuntimeSequenceOp>(referencedDevice, name);
  }
  return tables->lookupIn<AIE::RuntimeSequenceOp>(referencedDevice, name);
}
 
//===----------------------------------------------------------------------===//
// 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();
}
 
//===----------------------------------------------------------------------===//
// NpuWriteConfigOp
//===----------------------------------------------------------------------===//
 
static bool writesConfig(Operation *op) {
  return isa<AIEX::NpuLoadPdiOp, AIEX::NpuWriteConfigOp,
             AIEX::NpuWriteTeardownOp>(op);
}
 
// Visits the load_pdi, write_config and write_teardown ops before `op` in its
// block, nearest first, then those after it from the block's end, as the host
// runs the sequence again on each dispatch, then `op` itself. Returns the first
// for which `visit` returns true.
static Operation *findConfigOpBefore(Operation *op,
                                     function_ref<bool(Operation *)> visit) {
  for (Operation *prev = op->getPrevNode(); prev; prev = prev->getPrevNode())
    if (writesConfig(prev) && visit(prev))
      return prev;
  for (Operation &next : llvm::reverse(*op->getBlock())) {
    if (&next == op)
      break;
    if (writesConfig(&next) && visit(&next))
      return &next;
  }
  return visit(op) ? op : nullptr;
}
 
// The load_pdi that resets the array for `op`, or the write_teardown that
// leaves it so, or null.
static Operation *precedingReset(Operation *op) {
  return findConfigOpBefore(op, [&](Operation *prev) {
    return prev != op &&
           isa<AIEX::NpuLoadPdiOp, AIEX::NpuWriteTeardownOp>(prev);
  });
}
 
static bool loadsSamePdi(AIEX::NpuLoadPdiOp a, AIEX::NpuLoadPdiOp b) {
  return a.getDeviceRefAttr() == b.getDeviceRefAttr() &&
         a.getId() == b.getId() && a.getSize() == b.getSize() &&
         a.getAddress() == b.getAddress();
}
 
LogicalResult AIEX::NpuWriteConfigOp::verify() {
  Operation *found = precedingReset(*this);
  if (!found)
    return emitOpError("must follow the load_pdi that resets the array");
  auto reset = dyn_cast<AIEX::NpuLoadPdiOp>(found);
  if (!reset)
    return success();
  // The firmware skips the load if it loaded the same PDI last. That leaves
  // the array as a reset does only if every configuration since is torn down.
  bool tornDown = true;
  Operation *later = reset;
  auto before =
      cast<AIEX::NpuLoadPdiOp>(findConfigOpBefore(reset, [&](Operation *op) {
        if (isa<AIEX::NpuWriteConfigOp>(op) &&
            !isa<AIEX::NpuWriteTeardownOp>(later))
          tornDown = false;
        later = op;
        return isa<AIEX::NpuLoadPdiOp>(op);
      }));
  if (!loadsSamePdi(before, reset) || tornDown)
    return success();
  InFlightDiagnostic diag = emitOpError(
      "follows a load_pdi the firmware skips, as it loads the PDI loaded "
      "just before it, so the array is not reset");
  diag.attachNote(before.getLoc()) << "the PDI is loaded here";
  return diag;
}
 
LogicalResult
AIEX::NpuLoadPdiOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  FlatSymbolRefAttr ref = getDeviceRefAttr();
  if (!ref)
    return success();
  auto device = (*this)->getParentOfType<AIE::DeviceOp>();
  Operation *scope = device ? device->getParentOp()
                            : SymbolTable::getNearestSymbolTable(*this);
  if (!isa_and_present<AIE::DeviceOp>(symbolTable.lookupSymbolIn(scope, ref)))
    return emitOpError() << ref << " is not a device";
  return success();
}
 
LogicalResult
AIEX::NpuWriteConfigOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  Operation *scope = (*this)->getParentOfType<AIE::DeviceOp>()->getParentOp();
  if (!isa_and_present<AIE::DeviceOp>(
          symbolTable.lookupSymbolIn(scope, getDeviceRefAttr())))
    return emitOpError() << getDeviceRefAttr() << " is not a device";
  auto reset = dyn_cast_or_null<AIEX::NpuLoadPdiOp>(precedingReset(*this));
  FlatSymbolRefAttr resetRef = reset ? reset.getDeviceRefAttr() : nullptr;
  auto empty = resetRef ? dyn_cast_or_null<AIE::DeviceOp>(
                              symbolTable.lookupSymbolIn(scope, resetRef))
                        : nullptr;
  // The transaction conversion gives every device a `@configure` sequence; an
  // empty one still configures nothing.
  if (reset &&
      (!empty ||
       !llvm::all_of(empty.getBody()->without_terminator(), [](Operation &op) {
         return isa<AIE::RuntimeSequenceOp>(op) &&
                llvm::all_of(op.getRegion(0),
                             [](Block &block) { return block.empty(); });
       })))
    return emitOpError("must follow a load_pdi of an empty device, which "
                       "resets the array");
  return success();
}
 
//===----------------------------------------------------------------------===//
// NpuWriteTeardownOp
//===----------------------------------------------------------------------===//
 
LogicalResult AIEX::NpuWriteTeardownOp::verify() {
  Operation *prev = (*this)->getPrevNode();
  while (prev && !writesConfig(prev))
    prev = prev->getPrevNode();
  auto config = dyn_cast_or_null<AIEX::NpuWriteConfigOp>(prev);
  if (!config || config.getDeviceRefAttr() != getDeviceRefAttr())
    return emitOpError("must follow the write_config of ")
           << getDeviceRefAttr() << ", with no load_pdi between";
  return success();
}
 
LogicalResult
AIEX::NpuWriteTeardownOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  Operation *scope = (*this)->getParentOfType<AIE::DeviceOp>()->getParentOp();
  auto device = dyn_cast_or_null<AIE::DeviceOp>(
      symbolTable.lookupSymbolIn(scope, getDeviceRefAttr()));
  if (!device)
    return emitOpError() << getDeviceRefAttr() << " is not a device";
  Operation *shimStart = AIEX::findShimDmaStart(device);
  if (!shimStart)
    return success();
  InFlightDiagnostic diag =
      emitOpError() << getDeviceRefAttr()
                    << " starts a shim DMA channel, which teardown cannot stop";
  diag.attachNote(shimStart->getLoc()) << "started here";
  return diag;
}
 
Operation *AIEX::findShimDmaStart(AIE::DeviceOp device) {
  Operation *found = nullptr;
  device.walk([&](Operation *op) {
    if (!isa<AIE::DMAStartOp, AIE::DMAOp>(op) ||
        !isa<AIE::ShimDMAOp>(op->getParentOp()))
      return WalkResult::advance();
    found = op;
    return WalkResult::interrupt();
  });
  return found;
}