#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/Block.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Support/LLVM.h"

#include "llvm/ADT/DenseSet.h"
#include "llvm/Support/LogicalResult.h"

#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"

using namespace mlir;

namespace onnx_mlir {
namespace spatial {

namespace {

static FailureOr<ArrayRef<int64_t>> getWeightShapeForWeightedOp(Value weight) {
  auto shapedType = dyn_cast<ShapedType>(weight.getType());
  if (!shapedType)
    return failure();
  return shapedType.getShape();
}

template <typename ComputeBatchOpTy>
static bool isBatchOutputArgument(ComputeBatchOpTy batchOp, Value value) {
  if (batchOp.getNumResults() == 0)
    return false;
  auto blockArg = dyn_cast<BlockArgument>(value);
  if (!blockArg || blockArg.getOwner() != &batchOp.getBody().front())
    return false;

  unsigned argNumber = blockArg.getArgNumber();
  auto firstOutputArg = batchOp.getOutputArgument(0);
  if (!firstOutputArg)
    return false;
  unsigned firstOutputArgNumber = firstOutputArg->getArgNumber();
  return argNumber >= firstOutputArgNumber && argNumber < firstOutputArgNumber + batchOp.getNumResults();
}

template <typename ComputeOpTy>
static LogicalResult verifyStaticWeights(ComputeOpTy computeOp, StringRef kind) {
  for (Value weight : computeOp.getWeights())
    if (!isCompileTimeComputable(weight))
      return computeOp.emitOpError() << kind << " weights must be statically computed from constants";
  return success();
}

static bool isStaticScfForInductionVar(Value value) {
  auto blockArg = dyn_cast<BlockArgument>(value);
  if (!blockArg)
    return false;

  auto loop = dyn_cast_or_null<scf::ForOp>(blockArg.getOwner()->getParentOp());
  if (!loop || loop.getInductionVar() != value)
    return false;

  std::optional<int64_t> lowerBound = matchConstantIndexValue(loop.getLowerBound());
  std::optional<int64_t> upperBound = matchConstantIndexValue(loop.getUpperBound());
  std::optional<int64_t> step = matchConstantIndexValue(loop.getStep());
  return lowerBound && upperBound && step && *step > 0 && *upperBound >= *lowerBound;
}

static bool isStaticIndexExpr(Value value) {
  if (matchConstantIndexValue(value))
    return true;

  auto affineApply = value.getDefiningOp<affine::AffineApplyOp>();
  if (affineApply) {
    if (!isSingleResultSymbolFreeAffineMap(affineApply.getAffineMap()))
      return false;
    return llvm::all_of(affineApply.getMapOperands(), isStaticIndexExpr);
  }

  if (auto addOp = value.getDefiningOp<arith::AddIOp>())
    return isStaticIndexExpr(addOp.getLhs()) && isStaticIndexExpr(addOp.getRhs());

  if (auto mulOp = value.getDefiningOp<arith::MulIOp>())
    return isStaticIndexExpr(mulOp.getLhs()) && isStaticIndexExpr(mulOp.getRhs());

  return false;
}

static bool isSupportedLaneOffsetExpr(Value value, BlockArgument laneArg) {
  if (value == laneArg || isStaticIndexExpr(value) || isStaticScfForInductionVar(value))
    return true;

  auto affineApply = value.getDefiningOp<affine::AffineApplyOp>();
  if (affineApply) {
    if (!isSingleResultSymbolFreeAffineMap(affineApply.getAffineMap()))
      return false;
    if (!llvm::all_of(affineApply.getMapOperands(),
                      [&](Value operand) { return isSupportedLaneOffsetExpr(operand, laneArg); })) {
      return false;
    }
    return isDimAndConstantAffineExpr(affineApply.getAffineMap().getResult(0));
  }

  auto extractOp = value.getDefiningOp<tensor::ExtractOp>();
  if (extractOp) {
    auto constantTensor = extractOp.getTensor().getDefiningOp<arith::ConstantOp>();
    auto denseAttr = constantTensor ? dyn_cast<DenseIntElementsAttr>(constantTensor.getValue()) : nullptr;
    if (!denseAttr || denseAttr.getType().getRank() != 1 || extractOp.getIndices().size() != 1)
      return false;
    return isSupportedLaneOffsetExpr(extractOp.getIndices().front(), laneArg);
  }

  auto addOp = value.getDefiningOp<arith::AddIOp>();
  if (addOp)
    return (isSupportedLaneOffsetExpr(addOp.getLhs(), laneArg) && isStaticIndexExpr(addOp.getRhs()))
        || (isSupportedLaneOffsetExpr(addOp.getRhs(), laneArg) && isStaticIndexExpr(addOp.getLhs()));

  auto mulOp = value.getDefiningOp<arith::MulIOp>();
  if (!mulOp)
    return false;
  return (isSupportedLaneOffsetExpr(mulOp.getLhs(), laneArg) && isStaticIndexExpr(mulOp.getRhs()))
      || (isSupportedLaneOffsetExpr(mulOp.getRhs(), laneArg) && isStaticIndexExpr(mulOp.getLhs()));
}

static LogicalResult
verifyStaticUnitStrideExtractSliceOp(tensor::ExtractSliceOp sliceOp, BlockArgument laneArg, StringRef kind) {
  auto sourceType = dyn_cast<RankedTensorType>(sliceOp.getSource().getType());
  auto resultType = dyn_cast<RankedTensorType>(sliceOp.getResult().getType());
  if (!sourceType || !resultType || !sourceType.hasStaticShape() || !resultType.hasStaticShape())
    return sliceOp.emitOpError() << kind << " requires static ranked tensor types";
  if (!sliceOp.hasUnitStride())
    return sliceOp.emitOpError() << kind << " requires unit strides";

  for (int64_t size : sliceOp.getStaticSizes())
    if (ShapedType::isDynamic(size))
      return sliceOp.emitOpError() << kind << " requires static slice sizes";

  auto offsets = sliceOp.getOffsets();
  for (Value offset : offsets)
    if (!isSupportedLaneOffsetExpr(offset, laneArg))
      return sliceOp.emitOpError() << kind << " requires simple lane-dependent offsets";

  return success();
}

static LogicalResult verifyStaticUnitStrideParallelInsertSliceOp(tensor::ParallelInsertSliceOp sliceOp,
                                                                 BlockArgument laneArg,
                                                                 StringRef kind) {
  RankedTensorType sourceType = sliceOp.getSourceType();
  RankedTensorType destType = sliceOp.getDestType();
  if (!sourceType.hasStaticShape() || !destType.hasStaticShape())
    return sliceOp.emitOpError() << kind << " requires static ranked tensor types";
  if (!sliceOp.hasUnitStride())
    return sliceOp.emitOpError() << kind << " requires unit strides";

  for (int64_t size : sliceOp.getStaticSizes())
    if (ShapedType::isDynamic(size))
      return sliceOp.emitOpError() << kind << " requires static slice sizes";

  auto offsets = sliceOp.getOffsets();
  for (Value offset : offsets)
    if (!isSupportedLaneOffsetExpr(offset, laneArg))
      return sliceOp.emitOpError() << kind << " requires simple lane-dependent offsets";

  return success();
}

static Region* getParentRegion(Value value) {
  if (auto blockArg = dyn_cast<BlockArgument>(value))
    return blockArg.getOwner()->getParent();
  if (Operation* definingOp = value.getDefiningOp())
    return definingOp->getParentRegion();
  return nullptr;
}

static bool isDefinedInsideRegion(Value value, Region& region) {
  Region* parentRegion = getParentRegion(value);
  return parentRegion && (&region == parentRegion || region.isAncestor(parentRegion));
}

static bool isConstantExternalValue(Value value) {
  Operation* definingOp = value.getDefiningOp();
  return definingOp && definingOp->hasTrait<OpTrait::ConstantLike>();
}

static bool isRecordedDeferredCommunicationSource(Operation* op, Value value) {
  auto transfer = dyn_cast<SpatDeferredCommunicationOp>(op);
  return transfer && llvm::is_contained(transfer.getSources(), value);
}

static LogicalResult verifyOnlyConstantExternalValues(Operation* ownerOp, Region& region, StringRef kind) {
  bool hasFailure = false;
  region.walk([&](Operation* op) {
    for (OpOperand& operand : op->getOpOperands()) {
      Value value = operand.get();
      if (isDefinedInsideRegion(value, region) || isConstantExternalValue(value)
          || isRecordedDeferredCommunicationSource(op, value))
        continue;

      InFlightDiagnostic diagnostic =
        ownerOp->emitOpError() << kind << " body may not capture external values";
      diagnostic.attachNote(op->getLoc())
        << "owner='" << ownerOp->getName() << "' nestedOp='" << op->getName() << "' operand#"
        << operand.getOperandNumber() << " type=" << value.getType()
        << " category=" << (isa<TensorType>(value.getType()) ? "tensor" : (value.getType().isIndex() ? "index"
                                                                                                        : "scalar"));
      if (Operation* definingOp = value.getDefiningOp())
        diagnostic.attachNote(definingOp->getLoc()) << "defining op is '" << definingOp->getName() << "'";
      else if (auto blockArg = dyn_cast<BlockArgument>(value))
        diagnostic.attachNote(blockArg.getOwner()->getParentOp()->getLoc())
          << "value is block argument #" << blockArg.getArgNumber() << " of '"
          << blockArg.getOwner()->getParentOp()->getName() << "'";
      hasFailure = true;
    }
  });
  return success(!hasFailure);
}

static LogicalResult verifyYieldTypes(Operation* op, Region& region, TypeRange resultTypes, StringRef kind) {
  if (region.empty())
    return op->emitOpError() << kind << " requires a body block";
  Block& block = region.front();
  auto yield = dyn_cast_or_null<SpatYieldOp>(block.getTerminator());
  if (!yield)
    return op->emitOpError() << kind << " body must terminate with spat.yield";
  if (yield.getOutputs().size() != resultTypes.size())
    return op->emitOpError() << kind << " yield operand count must match result count";
  for (auto [yieldType, resultType] : llvm::zip(yield.getOutputs().getTypes(), resultTypes))
    if (yieldType != resultType)
      return op->emitOpError() << kind << " yield operand types must match result types";
  return success();
}

static LogicalResult verifyRegionArguments(Operation* op, Region& region, ValueRange operands, StringRef kind) {
  if (region.empty())
    return op->emitOpError() << kind << " requires a body block";
  Block& block = region.front();
  if (block.getNumArguments() != operands.size())
    return op->emitOpError() << kind << " body argument count must match operand count";
  for (auto [arg, operand] : llvm::zip(block.getArguments(), operands))
    if (arg.getType() != operand.getType())
      return op->emitOpError() << kind << " body argument types must match operand types";
  return success();
}

template <typename ComputeBatchOpTy>
static LogicalResult verifyBatchBody(ComputeBatchOpTy batchOp, Block& block, bool verifyLaneSliceOffsets = true) {
  if (batchOp.getNumResults() == 0) {
    auto yieldOp = dyn_cast_or_null<SpatYieldOp>(block.getTerminator());
    if (!yieldOp)
      return batchOp.emitError("resultless compute_batch body must terminate with spat.yield");
    if (yieldOp.getNumOperands() != 0)
      return batchOp.emitError("resultless compute_batch body yield must be empty");
  }
  else if (!isa_and_nonnull<SpatInParallelOp>(block.getTerminator())) {
    return batchOp.emitError("resultful compute_batch body must terminate with spat.in_parallel");
  }

  auto laneArg = batchOp.getLaneArgument();
  if (!laneArg)
    return batchOp.emitError("compute_batch body must have a lane block argument");
  if (verifyLaneSliceOffsets)
    for (auto& bodyOp : block) {
      if (auto extractSlice = dyn_cast<tensor::ExtractSliceOp>(&bodyOp))
        if (failed(verifyStaticUnitStrideExtractSliceOp(extractSlice, *laneArg, "tensor.extract_slice")))
          return failure();
    }
  return success();
}

} // namespace

LogicalResult SpatVMMOp::verify() {
  auto matrixShapeOpt = getWeightShapeForWeightedOp(getWeight());
  if (failed(matrixShapeOpt))
    return emitError("weight must be a shaped value");
  auto matrixShape = *matrixShapeOpt;
  auto vectorShape = getInput().getType().getShape();
  auto outputShape = getOutput().getType().getShape();

  if (matrixShape.size() != 2 || vectorShape.size() != 2 || outputShape.size() != 2)
    return emitError("matrix, vector and output must have rank 2");

  int64_t N = matrixShape[0];
  int64_t M = matrixShape[1];
  if (N <= 0 || M <= 0)
    return emitError("matrix shape must be (N, M) with N > 0 and M > 0");

  int64_t vector1 = vectorShape[0];
  int64_t vectorN = vectorShape[1];
  if (vectorN != N || vector1 != 1)
    return emitError("vector shape must be (1, N)");

  int64_t output1 = outputShape[0];
  int64_t outputM = outputShape[1];
  if (outputM != M || output1 != 1)
    return emitError("output shape must be (1, M)");

  return success();
}

LogicalResult SpatVVDMulOp::verify() {
  auto lhsType = dyn_cast<ShapedType>(getLhs().getType());
  auto rhsType = dyn_cast<ShapedType>(getRhs().getType());
  auto outputType = dyn_cast<ShapedType>(getOutput().getType());
  if (!lhsType || !rhsType || !outputType)
    return emitError("lhs, rhs, and output must be shaped values");
  if (!lhsType.hasRank() || !rhsType.hasRank() || !outputType.hasRank())
    return emitError("lhs, rhs, and output must have ranked types");

  ArrayRef<int64_t> lhsShape = lhsType.getShape();
  ArrayRef<int64_t> rhsShape = rhsType.getShape();
  ArrayRef<int64_t> outputShape = outputType.getShape();
  if (lhsShape.size() != 2 || rhsShape.size() != 2 || outputShape.size() != 2)
    return emitError("lhs, rhs, and output must have rank 2");
  if (lhsType.getElementType() != rhsType.getElementType() || lhsType.getElementType() != outputType.getElementType())
    return emitError("lhs, rhs, and output must have the same element type");
  if (lhsShape != rhsShape)
    return emitError("lhs and rhs vector shapes must match");
  if (lhsShape[0] != 1 || lhsShape[1] <= 0)
    return emitError("lhs and rhs vector shape must be (1, N) with N > 0");
  if (outputShape[0] != 1 || outputShape[1] != 1)
    return emitError("output shape must be (1, 1)");

  return success();
}

LogicalResult SpatVAddOp::verify() {
  if (failed(OpTrait::impl::verifyAtLeastNOperands(*this, 2)))
    return failure();
  return OpTrait::impl::verifySameOperandsAndResultType(*this);
}

LogicalResult SpatVSubOp::verify() {
  if (failed(OpTrait::impl::verifyAtLeastNOperands(*this, 2)))
    return failure();
  return OpTrait::impl::verifySameOperandsAndResultType(*this);
}

LogicalResult SpatVMaxOp::verify() {
  if (failed(OpTrait::impl::verifyAtLeastNOperands(*this, 2)))
    return failure();
  return OpTrait::impl::verifySameOperandsAndResultType(*this);
}

LogicalResult SpatExtractRowsOp::verify() {
  auto inputType = dyn_cast<ShapedType>(getInput().getType());
  if (!inputType || !inputType.hasRank() || inputType.getRank() != 2)
    return emitError("input must be a rank-2 shaped type");

  int64_t numRows = inputType.getShape()[0];
  int64_t numCols = inputType.getShape()[1];
  Type elementType = inputType.getElementType();

  if (numRows >= 0 && static_cast<int64_t>(getNumResults()) != numRows)
    return emitError("number of outputs must match the number of input rows");

  for (Type output : getResultTypes()) {
    auto outputType = dyn_cast<ShapedType>(output);
    if (!outputType || !outputType.hasRank() || outputType.getRank() != 2)
      return emitError("outputs must all be rank-2 shaped types");
    if (outputType.getElementType() != elementType)
      return emitError("output element types must match input element type");
    auto outputShape = outputType.getShape();
    if (outputShape[0] != 1)
      return emitError("each output must have exactly one row");
    if (numCols >= 0 && outputShape[1] != numCols)
      return emitError("output column count must match input column count");
  }

  return success();
}

LogicalResult SpatConcatOp::verify() {
  if (getInputs().empty())
    return emitError("requires at least one input");

  auto outputType = dyn_cast<ShapedType>(getOutput().getType());
  if (!outputType || !outputType.hasRank())
    return emitError("output must be a ranked shaped type");

  int64_t axis = getAxis();
  int64_t rank = outputType.getRank();
  if (axis < 0 || axis >= rank)
    return emitError("axis must be within the output rank");

  int64_t concatenatedDimSize = 0;
  bool concatenatedDimDynamic = false;
  Type outputElementType = outputType.getElementType();

  for (Value input : getInputs()) {
    auto inputType = dyn_cast<ShapedType>(input.getType());
    if (!inputType || !inputType.hasRank())
      return emitError("inputs must be ranked shaped types");
    if (inputType.getRank() != rank)
      return emitError("all inputs must have the same rank as the output");
    if (inputType.getElementType() != outputElementType)
      return emitError("all inputs must have the same element type as the output");

    for (int64_t dim = 0; dim < rank; ++dim) {
      if (dim == axis)
        continue;
      int64_t inputDim = inputType.getDimSize(dim);
      int64_t outputDim = outputType.getDimSize(dim);
      if (!ShapedType::isDynamic(inputDim) && !ShapedType::isDynamic(outputDim) && inputDim != outputDim)
        return emitError("non-concatenated dimensions must match the output shape");
    }

    int64_t inputConcatDim = inputType.getDimSize(axis);
    if (ShapedType::isDynamic(inputConcatDim)) {
      concatenatedDimDynamic = true;
      continue;
    }
    concatenatedDimSize += inputConcatDim;
  }

  int64_t outputConcatDim = outputType.getDimSize(axis);
  if (!concatenatedDimDynamic && !ShapedType::isDynamic(outputConcatDim) && concatenatedDimSize != outputConcatDim)
    return emitError("output concatenated dimension must equal the sum of input sizes");

  return success();
}

static bool isKnownLogicalLayout(StringRef layout) { return layout == "nchw"; }

static bool isKnownPhysicalLayout(StringRef layout) {
  return layout == "dense_nchw" || layout == "nchw_row_strip" || layout == "fragmented";
}

static LogicalResult verifyPlanTensorTypes(Operation* op, Value input, Value output, StringRef kind) {
  auto inputType = dyn_cast<RankedTensorType>(input.getType());
  auto outputType = dyn_cast<RankedTensorType>(output.getType());
  if (!inputType || !outputType)
    return op->emitOpError() << kind << " requires ranked tensor input and output types";
  if (inputType.getElementType() != outputType.getElementType())
    return op->emitOpError() << kind << " requires matching input/output element types";
  return success();
}

LogicalResult SpatConv2DPlanOp::verify() {
  auto inputType = dyn_cast<RankedTensorType>(getInput().getType());
  auto weightType = dyn_cast<RankedTensorType>(getWeight().getType());
  auto outputType = dyn_cast<RankedTensorType>(getOutput().getType());
  if (!inputType || !weightType || !outputType)
    return emitError("requires ranked tensor input, weight, and output");
  if (inputType.getRank() != 4 || weightType.getRank() != 4 || outputType.getRank() != 4)
    return emitError("requires rank-4 input, weight, and output tensors");
  if (!isKnownLogicalLayout(getLogicalLayout()))
    return emitError("requires a known logical layout");
  if (getPads().size() != 4)
    return emitError("requires exactly four pad values");
  if (getStrides().size() != 2)
    return emitError("requires exactly two stride values");
  if (getDilations().size() != 2)
    return emitError("requires exactly two dilation values");
  if (getGroup() < 1)
    return emitError("requires group >= 1");
  if (inputType.getElementType() != weightType.getElementType()
      || inputType.getElementType() != outputType.getElementType()) {
    return emitError("requires matching input, weight, and output element types");
  }
  if (getBias()) {
    auto biasType = dyn_cast<RankedTensorType>(getBias().getType());
    if (!biasType)
      return emitError("requires ranked tensor bias type");
    if (biasType.getElementType() != outputType.getElementType())
      return emitError("requires bias element type to match output element type");
  }
  return success();
}

LogicalResult SpatReluPlanOp::verify() {
  if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.relu_plan")))
    return failure();
  if (!isKnownLogicalLayout(getLogicalLayout()))
    return emitError("requires a known logical layout");
  return success();
}

LogicalResult SpatBiasAddPlanOp::verify() {
  if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.bias_add_plan")))
    return failure();
  if (!isKnownLogicalLayout(getLogicalLayout()))
    return emitError("requires a known logical layout");

  auto inputType = dyn_cast<RankedTensorType>(getInput().getType());
  auto biasType = dyn_cast<RankedTensorType>(getBias().getType());
  auto outputType = dyn_cast<RankedTensorType>(getOutput().getType());
  if (!inputType || !biasType || !outputType)
    return emitError("requires ranked tensor input, bias, and output");
  if (!inputType.hasStaticShape() || !biasType.hasStaticShape() || !outputType.hasStaticShape())
    return emitError("requires static tensor input, bias, and output");
  if (inputType != outputType)
    return emitError("requires matching input and output tensor types");
  if (outputType.getRank() != 4)
    return emitError("requires rank-4 input/output tensors");
  if (getLogicalLayout() != "nchw")
    return emitError("requires logical layout \"nchw\"");
  if (biasType.getElementType() != outputType.getElementType())
    return emitError("requires bias element type to match the output element type");

  ArrayRef<int64_t> biasShape = biasType.getShape();
  const int64_t channels = outputType.getDimSize(1);
  const bool supported = biasShape.empty() || (biasShape.size() == 1 && biasShape[0] == channels)
                         || (biasShape.size() == 2 && biasShape[0] == 1 && biasShape[1] == channels)
                         || (biasShape.size() == 4 && biasShape[0] == 1 && biasShape[1] == channels
                             && biasShape[2] == 1 && biasShape[3] == 1);
  if (!supported)
    return emitError("requires scalar or per-channel bias broadcastable over NCHW");
  return success();
}

LogicalResult SpatBlueprintOp::verify() {
  auto modeAttr = getModeAttr();
  bool isFragmentAssembly = modeAttr && modeAttr.getValue() == "fragment_assembly";
  if (!isFragmentAssembly && failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.blueprint")))
    return failure();
  if (!isKnownLogicalLayout(getLogicalLayout()))
    return emitError("requires a known logical layout");
  if (!isKnownPhysicalLayout(getPhysicalLayout()))
    return emitError("requires a known physical layout");

  auto logicalType = dyn_cast<RankedTensorType>(getOutput().getType());
  if (!logicalType)
    return emitError("requires ranked tensor output");

  auto offsets = getFragmentOffsets();
  auto sizes = getFragmentSizes();
  if (offsets.size() != sizes.size())
    return emitError("fragment offset and size arrays must have the same length");
  int64_t rank = logicalType.getRank();
  if (offsets.empty())
    return success();
  if (rank <= 0 || offsets.size() % rank != 0)
    return emitError("fragment metadata must be a whole number of rank-sized fragments");

  auto verifyBoundsOnly = [&](ArrayRef<int64_t> strideValues) -> LogicalResult {
    ArrayRef<int64_t> shape = logicalType.getShape();
    for (int64_t index = 0; index < static_cast<int64_t>(offsets.size()); ++index) {
      int64_t dim = index % rank;
      int64_t offset = offsets[index];
      int64_t size = sizes[index];
      int64_t stride = strideValues.empty() ? 1 : strideValues[index];
      if (offset < 0 || size < 0 || stride < 0)
        return emitError("fragment offsets, sizes, and strides must be non-negative");
      int64_t logicalDim = shape[dim];
      if (!ShapedType::isDynamic(logicalDim) && offset + size > logicalDim)
        return emitError("fragment bounds must stay within the logical tensor shape");
      if (stride != 1)
        return emitError("fragment assembly currently requires unit strides");
    }
    return success();
  };

  if (!isFragmentAssembly) {
    if (failed(verifyBoundsOnly({})))
      return failure();
    if (!getFragments().empty())
      return emitError("legacy blueprint does not accept extra fragment operands");
    if (getFragmentSourceOffsetsAttr() || getFragmentStridesAttr() || getConflictPolicyAttr()
        || getCoveragePolicyAttr())
      return emitError("legacy blueprint does not accept fragment assembly attributes");
    return success();
  }

  auto stridesAttr = getFragmentStridesAttr();
  auto operandIndicesAttr = getFragmentOperandIndicesAttr();
  auto sourceSlotsAttr = getFragmentSourceSlotsAttr();
  auto sourceOffsetsAttr = getFragmentSourceOffsetsAttr();
  if (!operandIndicesAttr)
    return emitError("fragment assembly blueprint requires fragment operand indices");
  if (!sourceSlotsAttr)
    return emitError("fragment assembly blueprint requires physical fragment source slots");
  if (!sourceOffsetsAttr)
    return emitError("fragment assembly blueprint requires fragment source offsets");
  if (!stridesAttr)
    return emitError("fragment assembly blueprint requires fragment strides");
  ArrayRef<int64_t> operandIndices = operandIndicesAttr.asArrayRef();
  ArrayRef<int64_t> sourceSlots = sourceSlotsAttr.asArrayRef();
  ArrayRef<int64_t> sourceOffsets = sourceOffsetsAttr.asArrayRef();
  ArrayRef<int64_t> strides = stridesAttr.asArrayRef();
  if (strides.size() != offsets.size())
    return emitError("fragment stride and offset arrays must have the same length");
  if (sourceOffsets.size() != operandIndices.size())
    return emitError("fragment source offset count must match fragment operand index count");
  if (sourceSlots.size() != operandIndices.size())
    return emitError("fragment source slot count must match fragment operand index count");
  if (!getConflictPolicyAttr() || !getCoveragePolicyAttr())
    return emitError("fragment assembly blueprint requires conflict and coverage policies");
  if (getConflictPolicy() != "disjoint")
    return emitError("fragment assembly blueprint currently supports only conflict_policy=\"disjoint\"");
  if (getCoveragePolicy() != "complete" && getCoveragePolicy() != "partial")
    return emitError("fragment assembly blueprint coverage_policy must be \"complete\" or \"partial\"");

  SmallVector<Value> operands;
  operands.push_back(getInput());
  llvm::append_range(operands, getFragments());
  int64_t operandCount = static_cast<int64_t>(operands.size());
  int64_t fragmentCount = static_cast<int64_t>(operandIndices.size());
  if (operandCount == 0)
    return emitError("fragment assembly blueprint requires at least one operand");
  if (static_cast<int64_t>(offsets.size()) != fragmentCount * rank)
    return emitError("fragment assembly metadata count must match operand count * result rank");
  if (failed(verifyBoundsOnly(strides)))
    return failure();

  SmallVector<std::pair<SmallVector<int64_t, 4>, SmallVector<int64_t, 4>>, 8> slices;
  slices.reserve(static_cast<size_t>(fragmentCount));
  SmallVector<int64_t, 8> fragmentCountsByOperand(static_cast<size_t>(operandCount), 0);
  auto expandFlatElementIndex = [](int64_t flatIndex, ArrayRef<int64_t> shape) {
    SmallVector<int64_t, 4> indices(shape.size(), 0);
    for (int64_t dim = static_cast<int64_t>(shape.size()) - 1; dim >= 0; --dim) {
      indices[dim] = flatIndex % shape[dim];
      flatIndex /= shape[dim];
    }
    return indices;
  };
  for (int64_t fragmentIndex = 0; fragmentIndex < fragmentCount; ++fragmentIndex) {
    int64_t operandIndex = operandIndices[fragmentIndex];
    if (operandIndex < 0 || operandIndex >= operandCount)
      return emitError("fragment assembly operand index is out of range");
    if (sourceSlots[fragmentIndex] < 0)
      return emitError("fragment assembly physical source slot must be nonnegative");
    if (sourceOffsets[fragmentIndex] < 0)
      return emitError("fragment assembly source offsets must be nonnegative");

    auto operandType = dyn_cast<RankedTensorType>(operands[operandIndex].getType());
    if (!operandType || !operandType.hasStaticShape())
      return emitError("fragment assembly blueprint requires static ranked tensor operands");
    if (operandType.getRank() != rank + 1)
      return emitError("fragment assembly physical operand must have one leading source-slot dimension");
    if (sourceSlots[fragmentIndex] >= operandType.getDimSize(0))
      return emitError("fragment assembly physical source slot is out of range");
    auto fragmentType = RankedTensorType::get(operandType.getShape().drop_front(), operandType.getElementType(), operandType.getEncoding());

    SmallVector<int64_t, 4> fragmentOffsets;
    SmallVector<int64_t, 4> fragmentSizes;
    fragmentOffsets.reserve(rank);
    fragmentSizes.reserve(rank);
    for (int64_t dim = 0; dim < rank; ++dim) {
      int64_t flatIndex = fragmentIndex * rank + dim;
      fragmentOffsets.push_back(offsets[flatIndex]);
      fragmentSizes.push_back(sizes[flatIndex]);
    }

    ++fragmentCountsByOperand[static_cast<size_t>(operandIndex)];
    int64_t fragmentElements = 1;
    for (int64_t dim = 0; dim < rank; ++dim)
      fragmentElements *= fragmentSizes[dim];
    if (sourceOffsets[fragmentIndex] + fragmentElements > fragmentType.getNumElements())
      return emitError("fragment assembly source offset exceeds the selected physical fragment bounds");
    SmallVector<int64_t, 4> sourceSliceOffsets =
      expandFlatElementIndex(sourceOffsets[fragmentIndex], fragmentType.getShape());
    for (int64_t dim = 0; dim < rank; ++dim)
      if (sourceSliceOffsets[dim] + fragmentSizes[dim] > fragmentType.getDimSize(dim))
        return emitError("fragment assembly source offset must describe a valid unit-stride slice");

    for (const auto& [existingOffsets, existingSizes] : slices) {
      bool overlaps = true;
      for (int64_t dim = 0; dim < rank; ++dim) {
        int64_t begin = fragmentOffsets[dim];
        int64_t end = begin + fragmentSizes[dim];
        int64_t existingBegin = existingOffsets[dim];
        int64_t existingEnd = existingBegin + existingSizes[dim];
        if (end <= existingBegin || existingEnd <= begin) {
          overlaps = false;
          break;
        }
      }
      if (overlaps)
        return emitError("fragment assembly blueprint requires disjoint static slices");
    }
    slices.push_back({std::move(fragmentOffsets), std::move(fragmentSizes)});
  }

  for (int64_t operandIndex = 0; operandIndex < operandCount; ++operandIndex) {
    if (fragmentCountsByOperand[static_cast<size_t>(operandIndex)] == 0)
      return emitError("fragment assembly blueprint requires every operand to contribute at least one fragment");
  }

  if (getCoveragePolicy() == "complete") {
    int64_t covered = 0;
    int64_t logicalElements = 1;
    for (int64_t dimSize : logicalType.getShape()) {
      if (ShapedType::isDynamic(dimSize))
        return emitError("fragment assembly complete coverage requires static result shape");
      logicalElements *= dimSize;
    }
    for (const auto& [ignoredOffsets, fragmentSizes] : slices) {
      int64_t fragmentElements = 1;
      for (int64_t dimSize : fragmentSizes)
        fragmentElements *= dimSize;
      covered += fragmentElements;
    }
    if (covered != logicalElements)
      return emitError("fragment assembly complete coverage must cover the whole result exactly");
  }

  return success();
}

LogicalResult SpatMaterializeLayoutOp::verify() {
  if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.materialize_layout")))
    return failure();
  if (!isKnownLogicalLayout(getLogicalLayout()))
    return emitError("requires a known logical layout");
  if (!isKnownPhysicalLayout(getSourcePhysicalLayout()))
    return emitError("requires a known source physical layout");
  if (!isKnownPhysicalLayout(getTargetPhysicalLayout()))
    return emitError("requires a known target physical layout");
  return success();
}

LogicalResult verifyComputeResultsUses(Operation* op) {
  if (!isAnySpatialComputeLike(op))
    return op->emitError("verifyComputeResultUses: op is not a Spatial compute-like operation");
  if (!llvm::all_of(op->getResults(), [](Value result) {
        return llvm::all_of(result.getUsers(), [result](Operation* op) {
          if (isRecordedDeferredCommunicationSource(op, result))
            return true;
          return !isAnySpatialComputeLike(op->getParentOp());
        });
      })) {
    return op->emitError("compute result used directly inside another Spatial compute body");
  }
  return success();
}

template <typename ComputeOpTy>
LogicalResult verifyComputeLikeOp(ComputeOpTy compute, StringRef opName) {
  unsigned expectedArgCount = compute.getWeights().size() + compute.getInputs().size();
  bool isScheduled = isa<SpatScheduledCompute>(compute.getOperation());
  if (compute.getBody().empty())
    return compute.emitOpError("compute body must have at least one block");

  SmallVector<Type> yieldedTypes;
  for (Block& block : compute.getBody()) {
    if ((!isScheduled && block.getNumArguments() != expectedArgCount)
        || (isScheduled && block.getNumArguments() < expectedArgCount))
      return compute.emitOpError("compute body must have weight and input block arguments");

    for (auto [weightIndex, weight] : llvm::enumerate(compute.getWeights()))
      if (block.getArgument(weightIndex).getType() != weight.getType())
        return compute.emitOpError("compute weight block argument types must match weight operand types exactly");
    for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs()))
      if (block.getArgument(compute.getWeights().size() + inputIndex).getType() != input.getType())
        return compute.emitOpError("compute input block argument types must match input operand types exactly");

    Operation* terminator = block.getTerminator();
    if (auto yieldOp = dyn_cast_or_null<SpatYieldOp>(terminator)) {
      auto realized = compute->template getAttrOfType<BoolAttr>("scheduled.realized");
      if (isScheduled && (!realized || !realized.getValue() || !compute.getBody().hasOneBlock()))
        return compute.emitOpError("scheduled compute blocks must terminate with spat.block_yield");
      llvm::append_range(yieldedTypes, yieldOp->getOperandTypes());
      continue;
    }
    auto blockYield = dyn_cast_or_null<SpatBlockYieldOp>(terminator);
    if (!blockYield || !isScheduled)
      return compute.emitOpError("ComputeOp must have a single yield operation");
    if (blockYield->getNumSuccessors() == 0)
      llvm::append_range(yieldedTypes, blockYield->getOperandTypes());
  }

  auto resultTypes = compute.getResultTypes();
  if (resultTypes.size() != yieldedTypes.size())
    return compute.emitOpError("ComputeOp must have same number of results as yielded operands");

  for (auto it : llvm::reverse(llvm::zip(resultTypes, yieldedTypes))) {
    auto resultType = std::get<0>(it);
    auto yieldType = std::get<1>(it);

    if (resultType != yieldType || failed(verifyCompatibleShape(resultType, yieldType)))
      return compute.emitOpError("ComputeOp output must be of the same type as yieldOp operand");

    if (auto resultRankedType = dyn_cast<RankedTensorType>(resultType)) {
      if (auto yieldRankedType = dyn_cast<RankedTensorType>(yieldType)) {
        if (resultRankedType.getEncoding() != yieldRankedType.getEncoding())
          return compute.emitOpError("ComputeOp output has an encoding while yieldOp operand does not have one");
      }
      else {
        return compute.emitOpError("ComputeOp output must have the same encoding as yieldOp operand");
      }
    }
    else if (dyn_cast<RankedTensorType>(yieldType)) {
      return compute.emitOpError("ComputeOp output must not have an encoding if yieldOp operand has one");
    }
  }

  if (compute.getBody().hasOneBlock())
    for (unsigned inputIndex = 0; inputIndex < compute.getInputs().size(); ++inputIndex)
      if (auto inputArg = compute.getInputArgument(inputIndex); !inputArg || inputArg->use_empty())
        return compute.emitOpError("ComputeOp block argument is not used");
  if (failed(verifyStaticWeights(compute, opName)))
    return failure();
  if (failed(verifyOnlyConstantExternalValues(compute.getOperation(), compute.getBody(), opName)))
    return failure();
  if (failed(verifyComputeResultsUses(compute.getOperation())))
    return failure();
  return success();
}

LogicalResult SpatGraphCompute::verify() { return verifyComputeLikeOp(*this, "spat.graph_compute"); }

LogicalResult SpatScheduledCompute::verify() { return verifyComputeLikeOp(*this, "spat.scheduled_compute"); }

LogicalResult SpatBlockYieldOp::verify() {
  if (getOperation()->getNumSuccessors() > 1)
    return emitOpError("may target at most one next scheduled block");
  Operation* parent = getOperation()->getParentOp();
  if (!isa_and_nonnull<SpatScheduledCompute>(parent))
    return emitOpError("expected spat.scheduled_compute parent");
  if (getOperation()->getNumSuccessors() == 1) {
    Block* next = getOperation()->getSuccessor(0);
    if (getOperation()->getNumOperands() != next->getNumArguments())
      return emitOpError("successor operand count must match next block argument count");
    for (auto [operand, argument] : llvm::zip(getOperation()->getOperands(), next->getArguments()))
      if (operand.getType() != argument.getType())
        return emitOpError("successor operand types must match next block argument types");
  }
  return success();
}

LogicalResult SpatDeferredCommunicationOp::verify() {
  if (getSources().empty())
    return emitOpError("requires at least one source");
  static constexpr StringLiteral staleAttributes[] = {
    "exchangeId", "logicalProducer", "logicalConsumer", "sourceClass", "targetClass", "sourceCore",
    "targetCore", "sourceLane", "targetLane", "transferKind", "resultIndex", "projectedTransfer",
    "hostOutputOwner", "source_cpus", "source_classes", "source_lane_ranges", "target_cpus",
    "target_classes", "target_lane_ranges", "batched", "source_operand_for_scheduled_lane",
    "multi_source_payload"};
  for (StringLiteral name : staleAttributes)
    if (getOperation()->hasAttr(name))
      return emitOpError() << "does not accept stale routing attribute '" << name
                           << "'; source selection and shaping belong in the body and routing is derived in Phase 2";
  if (failed(verifyRegionArguments(getOperation(), getBody(), getSources(), "spat.deferred_communication")))
    return failure();
  return verifyYieldTypes(getOperation(), getBody(), getOperation()->getResultTypes(), "spat.deferred_communication");
}

template <typename ComputeBatchOpTy>
LogicalResult verifyComputeBatchLikeOp(ComputeBatchOpTy batch, StringRef opName) {
  int32_t count = batch.getLaneCount();
  if (count <= 0)
    return batch.emitOpError("laneCount must be positive");

  auto laneCountSz = static_cast<size_t>(count);

  if (auto coreIdAttr = batch->getAttr(kCoreIdsAttrName)) {
    auto coreIdsAttr = dyn_cast<DenseI32ArrayAttr>(coreIdAttr);
    if (!coreIdsAttr)
      return batch.emitOpError("compute_batch coreIds attribute must be a dense i32 array");
    if (coreIdsAttr.size() != static_cast<int64_t>(laneCountSz))
      return batch.emitOpError("compute_batch coreIds array length must match laneCount");
    if (llvm::any_of(coreIdsAttr.asArrayRef(), [](int32_t coreId) { return coreId < 0; }))
      return batch.emitOpError("compute_batch coreIds values must be non-negative");
    DenseSet<int32_t> seenCoreIds;
    for (int32_t coreId : coreIdsAttr.asArrayRef())
      if (!seenCoreIds.insert(coreId).second)
        return batch.emitOpError("compute_batch coreIds values must be unique");
  }

  if (batch.getBody().empty())
    return batch.emitOpError("compute_batch body must have at least one block");

  unsigned expectedArgCount = 1 + batch.getWeights().size() + batch.getInputs().size() + batch.getNumResults();
  bool verifyLaneSliceOffsets = !isa<SpatScheduledComputeBatch>(batch.getOperation());
  for (Block& block : batch.getBody()) {
    if (block.getNumArguments() == 0)
      return batch.emitOpError("compute_batch body must have exactly one lane block argument");
    if (block.getNumArguments() != expectedArgCount)
      return batch.emitOpError(
        "compute_batch body block arguments must match lane, weight, input, and output operands/results");
    if (!block.getArgument(0).getType().isIndex())
      return batch.emitOpError("compute_batch first block argument must have index type");

    for (auto [weightIndex, weight] : llvm::enumerate(batch.getWeights()))
      if (block.getArgument(1 + weightIndex).getType() != weight.getType())
        return batch.emitOpError("compute_batch weight block argument types must match weight operand types exactly");
    for (auto [inputIndex, input] : llvm::enumerate(batch.getInputs()))
      if (block.getArgument(1 + batch.getWeights().size() + inputIndex).getType() != input.getType())
        return batch.emitOpError("compute_batch input block argument types must match input operand types exactly");
    for (auto [resultIndex, resultType] : llvm::enumerate(batch.getResultTypes()))
      if (block.getArgument(1 + batch.getWeights().size() + batch.getInputs().size() + resultIndex).getType()
          != resultType)
        return batch.emitOpError("compute_batch output block argument types must match result types exactly");

    if (failed(verifyBatchBody(batch, block, verifyLaneSliceOffsets)))
      return failure();
  }

  if (failed(verifyComputeResultsUses(batch.getOperation())))
    return failure();
  if (failed(verifyStaticWeights(batch, opName)))
    return failure();
  if (failed(verifyOnlyConstantExternalValues(batch.getOperation(), batch.getBody(), opName)))
    return failure();
  return success();
}

LogicalResult SpatGraphComputeBatch::verify() { return verifyComputeBatchLikeOp(*this, "spat.graph_compute_batch"); }

LogicalResult SpatScheduledComputeBatch::verify() {
  return verifyComputeBatchLikeOp(*this, "spat.scheduled_compute_batch");
}

LogicalResult SpatInParallelOp::verify() {
  Operation* parent = getOperation()->getParentOp();
  if (!isAnySpatialComputeBatchLike(parent))
    return emitOpError("expected spat.graph_compute_batch or spat.scheduled_compute_batch parent");
  if (parent->getNumResults() == 0)
    return emitOpError("requires a resultful spat.compute_batch parent");

  std::optional<BlockArgument> laneArg;
  if (auto graphBatch = dyn_cast<SpatGraphComputeBatch>(parent))
    laneArg = graphBatch.getLaneArgument();
  else
    laneArg = cast<SpatScheduledComputeBatch>(parent).getLaneArgument();
  if (!laneArg)
    return emitOpError("expected compute_batch lane block argument");
  for (Operation& op : getRegion().front().getOperations()) {
    auto insertSliceOp = dyn_cast<tensor::ParallelInsertSliceOp>(&op);
    if (!insertSliceOp)
      return emitOpError("expected only tensor.parallel_insert_slice ops");

    if (failed(verifyStaticUnitStrideParallelInsertSliceOp(insertSliceOp, *laneArg, "tensor.parallel_insert_slice")))
      return failure();

    MutableOperandRange destinations = insertSliceOp.getUpdatedDestinations();
    for (OpOperand& destination : destinations)
      if ((isa<SpatGraphComputeBatch>(parent)
             && !isBatchOutputArgument(cast<SpatGraphComputeBatch>(parent), destination.get()))
          || (isa<SpatScheduledComputeBatch>(parent)
              && !isBatchOutputArgument(cast<SpatScheduledComputeBatch>(parent), destination.get())))
        return op.emitOpError("may only insert into a compute_batch output block argument");
  }

  return success();
}

} // namespace spatial
} // namespace onnx_mlir