#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/PatternMatch.h"

#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"

#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Dialect/ONNX/ONNXOps.hpp"

using namespace mlir;

namespace onnx_mlir {
namespace {

static FailureOr<SmallVector<int64_t>> inferSupportedBatchShape(ArrayRef<int64_t> lhsBatchShape,
                                                                ArrayRef<int64_t> rhsBatchShape) {
  if (lhsBatchShape.empty())
    return SmallVector<int64_t>(rhsBatchShape.begin(), rhsBatchShape.end());
  if (rhsBatchShape.empty())
    return SmallVector<int64_t>(lhsBatchShape.begin(), lhsBatchShape.end());
  if (!llvm::equal(lhsBatchShape, rhsBatchShape))
    return failure();
  return SmallVector<int64_t>(lhsBatchShape.begin(), lhsBatchShape.end());
}

static Value
collapseBatchDims(Value value, int64_t batchSize, int64_t rows, int64_t cols, PatternRewriter& rewriter, Location loc) {
  auto type = cast<RankedTensorType>(value.getType());
  if (type.getRank() == 2 || type.getRank() == 3)
    return value;

  auto collapsedType = RankedTensorType::get({batchSize, rows, cols}, type.getElementType(), type.getEncoding());
  SmallVector<ReassociationIndices> reassociation = {ReassociationIndices {},
                                                     ReassociationIndices {static_cast<int64_t>(type.getRank() - 2)},
                                                     ReassociationIndices {static_cast<int64_t>(type.getRank() - 1)}};
  for (int64_t dim = 0; dim < type.getRank() - 2; ++dim)
    reassociation.front().push_back(dim);

  auto buildCollapsed = [&](Value input) -> Value {
    return tensor::CollapseShapeOp::create(rewriter, loc, collapsedType, input, reassociation);
  };
  return materializeOrComputeUnary(value, collapsedType, rewriter, loc, buildCollapsed);
}

static Value
expandBatchDims(Value value, RankedTensorType outputType, size_t batchRank, PatternRewriter& rewriter, Location loc) {
  if (cast<RankedTensorType>(value.getType()) == outputType)
    return value;

  SmallVector<ReassociationIndices> reassociation = {ReassociationIndices {},
                                                     ReassociationIndices {static_cast<int64_t>(batchRank)},
                                                     ReassociationIndices {static_cast<int64_t>(batchRank + 1)}};
  for (size_t dim = 0; dim < batchRank; ++dim)
    reassociation.front().push_back(static_cast<int64_t>(dim));

  auto buildExpanded = [&](Value input) -> Value {
    return tensor::ExpandShapeOp::create(rewriter, loc, outputType, input, reassociation).getResult();
  };
  return materializeOrComputeUnary(value, outputType, rewriter, loc, buildExpanded);
}

static Value ensureBatchedTensor(
  Value value, int64_t batchSize, int64_t rows, int64_t cols, PatternRewriter& rewriter, Location loc) {
  auto type = cast<RankedTensorType>(value.getType());
  if (type.getRank() == 3)
    return value;

  auto batchedType = RankedTensorType::get({batchSize, rows, cols}, type.getElementType(), type.getEncoding());
  auto buildExpanded = [&](Value input) -> Value {
    return tensor::ExpandShapeOp::create(rewriter,
                                         loc,
                                         batchedType,
                                         input,
                                         SmallVector<ReassociationIndices> {
                                           {0, 1},
                                           {2}
    });
  };
  return materializeOrComputeUnary(value, batchedType, rewriter, loc, buildExpanded);
}

static Value extractBatchMatrix(Value value,
                                int64_t batchIndex,
                                int64_t batchSize,
                                int64_t rows,
                                int64_t cols,
                                PatternRewriter& rewriter,
                                Location loc) {
  auto type = cast<RankedTensorType>(value.getType());
  if (type.getRank() == 2)
    return value;

  auto sliceType = RankedTensorType::get({1, rows, cols}, type.getElementType());
  SmallVector<OpFoldResult> offsets = {
    rewriter.getIndexAttr(batchSize == 1 ? 0 : batchIndex), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
  SmallVector<OpFoldResult> sizes = {
    rewriter.getIndexAttr(1), rewriter.getIndexAttr(rows), rewriter.getIndexAttr(cols)};
  SmallVector<OpFoldResult> strides = getUnitStrides(rewriter, 3);
  auto matrixType = RankedTensorType::get({rows, cols}, type.getElementType());
  auto buildMatrix = [&](Value input) -> Value {
    Value slice = tensor::ExtractSliceOp::create(rewriter, loc, sliceType, input, offsets, sizes, strides);
    return tensor::CollapseShapeOp::create(rewriter,
                                           loc,
                                           matrixType,
                                           slice,
                                           SmallVector<ReassociationIndices> {
                                             {0, 1},
                                             {2}
    });
  };

  return materializeOrComputeUnary(value, matrixType, rewriter, loc, buildMatrix);
}

static Value transposeLastTwoDims(Value value, PatternRewriter& rewriter, Location loc) {
  auto type = cast<RankedTensorType>(value.getType());
  auto shape = type.getShape();
  auto createONNXTranspose = [&](RankedTensorType resultType, ArrayRef<int64_t> permutation) {
    return ONNXTransposeOp::create(rewriter, loc, resultType, value, rewriter.getI64ArrayAttr(permutation)).getResult();
  };
  if (type.getRank() == 2) {
    auto resultType = RankedTensorType::get({shape[1], shape[0]}, type.getElementType(), type.getEncoding());
    return createONNXTranspose(resultType, {1, 0});
  }

  auto resultType = RankedTensorType::get({shape[0], shape[2], shape[1]}, type.getElementType(), type.getEncoding());
  return createONNXTranspose(resultType, {0, 2, 1});
}

static Value createZeroPaddedTensor(Value value, RankedTensorType resultType, PatternRewriter& rewriter, Location loc) {
  auto sourceType = cast<RankedTensorType>(value.getType());
  SmallVector<OpFoldResult> lowPads(sourceType.getRank(), rewriter.getIndexAttr(0));
  SmallVector<OpFoldResult> highPads;
  highPads.reserve(sourceType.getRank());
  for (auto [sourceDim, resultDim] : llvm::zip(sourceType.getShape(), resultType.getShape()))
    highPads.push_back(rewriter.getIndexAttr(resultDim - sourceDim));

  auto padOp = tensor::PadOp::create(rewriter, loc, resultType, value, lowPads, highPads);
  auto* padBlock = new Block();
  for (int64_t i = 0; i < sourceType.getRank(); ++i)
    padBlock->addArgument(rewriter.getIndexType(), loc);
  padOp.getRegion().push_back(padBlock);
  rewriter.setInsertionPointToStart(padBlock);
  auto zero = getOrCreateConstant(
    rewriter, padOp.getOperation(), rewriter.getZeroAttr(sourceType.getElementType()), sourceType.getElementType());
  tensor::YieldOp::create(rewriter, loc, zero);
  rewriter.setInsertionPointAfter(padOp);
  return padOp.getResult();
}

static Value createPaddedBatchedInputCompute(Value input,
                                             RankedTensorType paddedInputType,
                                             PatternRewriter& rewriter,
                                             Location loc) {
  auto inputType = cast<RankedTensorType>(input.getType());
  if (inputType == paddedInputType)
    return input;

  auto computeOp = createSpatCompute<1>(rewriter, loc, TypeRange {paddedInputType}, {}, input, [&](Value computeInput) {
    Value paddedInput = createZeroPaddedTensor(computeInput, paddedInputType, rewriter, loc);
    spatial::SpatYieldOp::create(rewriter, loc, paddedInput);
  });
  return computeOp.getResult(0);
}

static FailureOr<Value> materializePaddedBatchedWeight(
  Value value, int64_t sourceBatch, int64_t targetBatch, RankedTensorType resultType, PatternRewriter& rewriter) {
  auto sourceType = cast<RankedTensorType>(value.getType());
  if (sourceType == resultType)
    return value;

  auto denseAttr = getHostConstDenseElementsAttr(value);
  if (!denseAttr)
    return failure();

  const int64_t sourceRows = sourceType.getRank() == 2 ? sourceType.getDimSize(0) : sourceType.getDimSize(1);
  const int64_t sourceCols = sourceType.getRank() == 2 ? sourceType.getDimSize(1) : sourceType.getDimSize(2);
  const int64_t targetRows = resultType.getDimSize(1);
  const int64_t targetCols = resultType.getDimSize(2);
  SmallVector<Attribute> sourceValues(denseAttr.getValues<Attribute>());
  SmallVector<Attribute> resultValues(resultType.getNumElements(), rewriter.getZeroAttr(resultType.getElementType()));

  for (int64_t batchIdx = 0; batchIdx < targetBatch; ++batchIdx) {
    const int64_t sourceBatchIdx = sourceType.getRank() == 2 ? 0 : (sourceBatch == 1 ? 0 : batchIdx);
    const int64_t sourceBatchBase = sourceType.getRank() == 2 ? 0 : sourceBatchIdx * sourceRows * sourceCols;
    const int64_t targetBatchBase = batchIdx * targetRows * targetCols;
    for (int64_t row = 0; row < sourceRows; ++row)
      for (int64_t col = 0; col < sourceCols; ++col)
        resultValues[targetBatchBase + row * targetCols + col] = sourceValues[sourceBatchBase + row * sourceCols + col];
  }

  auto resultAttr = DenseElementsAttr::get(resultType, resultValues);
  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), resultAttr, resultType);
}

static Value extractBatchedATile(Value a,
                                 int64_t sourceBatchCount,
                                 Value batch,
                                 Value row,
                                 Value kOffset,
                                 RankedTensorType aTileType,
                                 PatternRewriter& rewriter,
                                 Location loc) {
  auto aSliceType = RankedTensorType::get({1, 1, aTileType.getDimSize(1)}, aTileType.getElementType());
  SmallVector<OpFoldResult> offsets {
    sourceBatchCount == 1 ? OpFoldResult(rewriter.getIndexAttr(0)) : OpFoldResult(batch), row, kOffset};
  SmallVector<OpFoldResult> sizes {
    rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(aTileType.getDimSize(1))};
  auto slice =
    tensor::ExtractSliceOp::create(rewriter, loc, aSliceType, a, offsets, sizes, getUnitStrides(rewriter, 3));
  return tensor::CollapseShapeOp::create(rewriter,
                                         loc,
                                         aTileType,
                                         slice,
                                         SmallVector<ReassociationIndices> {
                                           {0, 1},
                                           {2}
  });
}

static Value extractBatchedBTile(Value b,
                                 int64_t sourceBatchCount,
                                 Value batch,
                                 Value kOffset,
                                 Value hOffset,
                                 RankedTensorType bTileType,
                                 PatternRewriter& rewriter,
                                 Location loc) {
  auto bSliceType =
    RankedTensorType::get({1, bTileType.getDimSize(0), bTileType.getDimSize(1)}, bTileType.getElementType());
  SmallVector<OpFoldResult> offsets {
    sourceBatchCount == 1 ? OpFoldResult(rewriter.getIndexAttr(0)) : OpFoldResult(batch), kOffset, hOffset};
  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
                                   rewriter.getIndexAttr(bTileType.getDimSize(0)),
                                   rewriter.getIndexAttr(bTileType.getDimSize(1))};
  auto slice =
    tensor::ExtractSliceOp::create(rewriter, loc, bSliceType, b, offsets, sizes, getUnitStrides(rewriter, 3));
  return tensor::CollapseShapeOp::create(rewriter,
                                         loc,
                                         bTileType,
                                         slice,
                                         SmallVector<ReassociationIndices> {
                                           {0, 1},
                                           {2}
  });
}

static Value getBatchLaneIndex(
  Value lane, int64_t numOutRows, int64_t numKSlices, int64_t numOutHSlices, PatternRewriter& rewriter, Location loc) {
  return affineFloorDivConst(
    rewriter, loc, lane, numOutRows * numKSlices * numOutHSlices, rewriter.getInsertionBlock()->getParentOp());
}

static FailureOr<spatial::SpatComputeBatch> createBatchedVmmBatch(Value a,
                                                                  Value b,
                                                                  RankedTensorType aType,
                                                                  int64_t aBatchCount,
                                                                  RankedTensorType bType,
                                                                  int64_t bBatchCount,
                                                                  RankedTensorType partialPiecesType,
                                                                  int64_t numOutRows,
                                                                  int64_t numKSlices,
                                                                  int64_t numOutHSlices,
                                                                  PatternRewriter& rewriter,
                                                                  Location loc) {
  const int64_t laneCount = partialPiecesType.getDimSize(0);
  auto batchOp = createSpatComputeBatch(
    rewriter,
    loc,
    TypeRange {partialPiecesType},
    laneCount,
    ValueRange {b},
    ValueRange {a},
    [&](detail::SpatComputeBatchBodyArgs args) {
      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
      Value row = affineModConst(rewriter, loc, args.lane, numOutRows, anchorOp);
      Value outerLane = affineFloorDivConst(rewriter, loc, args.lane, numOutRows, anchorOp);
      Value batch = getBatchLaneIndex(args.lane, numOutRows, numKSlices, numOutHSlices, rewriter, loc);
      Value sliceLane = affineModConst(rewriter, loc, outerLane, numKSlices * numOutHSlices, anchorOp);
      Value kSlice = affineModConst(rewriter, loc, sliceLane, numKSlices, anchorOp);
      Value hSlice = affineFloorDivConst(rewriter, loc, sliceLane, numKSlices, anchorOp);
      Value kOffset = affineMulConst(rewriter, loc, kSlice, crossbarSize.getValue(), anchorOp);
      Value hOffset = affineMulConst(rewriter, loc, hSlice, crossbarSize.getValue(), anchorOp);

      auto aTileType =
        RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, aType.getElementType());
      auto bTileType = RankedTensorType::get(
        {static_cast<int64_t>(crossbarSize.getValue()), static_cast<int64_t>(crossbarSize.getValue())},
        bType.getElementType());
      auto pieceType =
        RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, partialPiecesType.getElementType());

      Value aTile =
        extractBatchedATile(args.inputs.front(), aBatchCount, batch, row, kOffset, aTileType, rewriter, loc);
      Value bTile =
        extractBatchedBTile(args.weights.front(), bBatchCount, batch, kOffset, hOffset, bTileType, rewriter, loc);
      Value piece = spatial::SpatVMMOp::create(rewriter, loc, pieceType, bTile, aTile).getResult();

      SmallVector<OpFoldResult> pieceOffsets {args.lane, rewriter.getIndexAttr(0)};
      SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
      createParallelInsertSliceIntoBatchOutput(
        rewriter, loc, piece, args.outputs.front(), pieceOffsets, pieceSizes, getUnitStrides(rewriter, 2));
    });
  if (failed(batchOp))
    return failure();
  return *batchOp;
}

static Value extractDynamicBatchedBColumn(Value matrix,
                                          int64_t sourceBatchCount,
                                          Value batch,
                                          Value column,
                                          RankedTensorType vectorType,
                                          PatternRewriter& rewriter,
                                          Location loc) {
  auto columnSliceType = RankedTensorType::get({1, vectorType.getDimSize(1), 1}, vectorType.getElementType());
  SmallVector<OpFoldResult> offsets {sourceBatchCount == 1 ? OpFoldResult(rewriter.getIndexAttr(0))
                                                           : OpFoldResult(batch),
                                     rewriter.getIndexAttr(0),
                                     column};
  SmallVector<OpFoldResult> sizes {
    rewriter.getIndexAttr(1), rewriter.getIndexAttr(vectorType.getDimSize(1)), rewriter.getIndexAttr(1)};
  SmallVector<OpFoldResult> strides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
  Value columnSlice = tensor::ExtractSliceOp::create(rewriter, loc, columnSliceType, matrix, offsets, sizes, strides);
  auto collapsedType = RankedTensorType::get({vectorType.getDimSize(1)}, vectorType.getElementType());
  Value collapsed = tensor::CollapseShapeOp::create(rewriter,
                                                    loc,
                                                    collapsedType,
                                                    columnSlice,
                                                    SmallVector<ReassociationIndices> {
                                                      {0, 1, 2}
  })
                      .getResult();
  return tensor::ExpandShapeOp::create(rewriter,
                                       loc,
                                       vectorType,
                                       collapsed,
                                       SmallVector<ReassociationIndices> {
                                         {0, 1}
  })
    .getResult();
}

static Value extractDynamicBatchedRowVector(Value matrix,
                                            int64_t sourceBatchCount,
                                            Value batch,
                                            Value row,
                                            RankedTensorType vectorType,
                                            PatternRewriter& rewriter,
                                            Location loc) {
  auto rowSliceType = RankedTensorType::get({1, 1, vectorType.getDimSize(1)}, vectorType.getElementType());
  SmallVector<OpFoldResult> offsets {sourceBatchCount == 1 ? OpFoldResult(rewriter.getIndexAttr(0))
                                                           : OpFoldResult(batch),
                                     row,
                                     rewriter.getIndexAttr(0)};
  SmallVector<OpFoldResult> sizes {
    rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(vectorType.getDimSize(1))};
  auto rowSlice =
    tensor::ExtractSliceOp::create(rewriter, loc, rowSliceType, matrix, offsets, sizes, getUnitStrides(rewriter, 3));
  return tensor::CollapseShapeOp::create(rewriter,
                                         loc,
                                         vectorType,
                                         rowSlice,
                                         SmallVector<ReassociationIndices> {
                                           {0, 1},
                                           {2}
  });
}

static FailureOr<spatial::SpatComputeBatch> createBatchedVvdmulBatch(Value a,
                                                                     int64_t aBatchCount,
                                                                     Value b,
                                                                     int64_t bBatchCount,
                                                                     RankedTensorType aType,
                                                                     RankedTensorType bType,
                                                                     RankedTensorType scalarPiecesType,
                                                                     RankedTensorType outType,
                                                                     PatternRewriter& rewriter,
                                                                     Location loc) {
  const int64_t numBatches = outType.getDimSize(0);
  const int64_t numOutRows = outType.getDimSize(1);
  const int64_t numOutCols = outType.getDimSize(2);
  const int64_t reductionSize = aType.getDimSize(2);
  const int64_t laneCount = numBatches * numOutRows * numOutCols;
  auto batchOp = createSpatComputeBatch(
    rewriter,
    loc,
    TypeRange {scalarPiecesType},
    laneCount,
    ValueRange {},
    ValueRange {a, b},
    [&](detail::SpatComputeBatchBodyArgs args) {
      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
      Value batch = affineFloorDivConst(rewriter, loc, args.lane, numOutRows * numOutCols, anchorOp);
      Value batchLane = affineModConst(rewriter, loc, args.lane, numOutRows * numOutCols, anchorOp);
      Value row = affineFloorDivConst(rewriter, loc, batchLane, numOutCols, anchorOp);
      Value column = affineModConst(rewriter, loc, batchLane, numOutCols, anchorOp);

      auto vectorType = RankedTensorType::get({1, reductionSize}, aType.getElementType());
      auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
      Value aVector =
        extractDynamicBatchedRowVector(args.inputs[0], aBatchCount, batch, row, vectorType, rewriter, loc);
      Value bVector =
        extractDynamicBatchedBColumn(args.inputs[1], bBatchCount, batch, column, vectorType, rewriter, loc);
      Value scalar = spatial::SpatVVDMulOp::create(rewriter, loc, scalarType, aVector, bVector).getResult();
      SmallVector<OpFoldResult> outputOffsets {args.lane, rewriter.getIndexAttr(0)};
      SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
      createParallelInsertSliceIntoBatchOutput(
        rewriter, loc, scalar, args.outputs.front(), outputOffsets, scalarSizes, getUnitStrides(rewriter, 2));
    });
  if (failed(batchOp))
    return failure();
  return *batchOp;
}

static FailureOr<Value> createBatchedDynamicOutputCompute(Value scalarPieces,
                                                          RankedTensorType scalarPiecesType,
                                                          RankedTensorType outType,
                                                          PatternRewriter& rewriter,
                                                          Location loc) {
  const int64_t laneCount = scalarPiecesType.getDimSize(0);
  const int64_t numOutRows = outType.getDimSize(1);
  const int64_t numOutCols = outType.getDimSize(2);
  auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
  auto outputScalarType = RankedTensorType::get({1, 1, 1}, outType.getElementType());

  auto computeOp = createSpatCompute<1>(
    rewriter, loc, TypeRange {outType}, {}, ValueRange {scalarPieces}, [&](Value pieces) -> LogicalResult {
      Value outputInit =
        tensor::EmptyOp::create(rewriter, loc, outType.getShape(), outType.getElementType()).getResult();
      Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
      Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1);
      Value cLaneCount = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), laneCount);
      auto loop = buildNormalizedScfFor(
        rewriter,
        loc,
        c0,
        cLaneCount,
        c1,
        ValueRange {outputInit},
        [&](OpBuilder&, Location nestedLoc, Value lane, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
          Value outputAcc = iterArgs.front();
          Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
          Value batch = affineFloorDivConst(rewriter, nestedLoc, lane, numOutRows * numOutCols, anchorOp);
          Value batchLane = affineModConst(rewriter, nestedLoc, lane, numOutRows * numOutCols, anchorOp);
          Value row = affineFloorDivConst(rewriter, nestedLoc, batchLane, numOutCols, anchorOp);
          Value column = affineModConst(rewriter, nestedLoc, batchLane, numOutCols, anchorOp);
          SmallVector<OpFoldResult> scalarOffsets {lane, rewriter.getIndexAttr(0)};
          SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
          Value scalar = tensor::ExtractSliceOp::create(
            rewriter, nestedLoc, scalarType, pieces, scalarOffsets, scalarSizes, getUnitStrides(rewriter, 2));
          Value expanded = tensor::ExpandShapeOp::create(rewriter,
                                                         nestedLoc,
                                                         outputScalarType,
                                                         scalar,
                                                         SmallVector<ReassociationIndices> {
                                                           {0},
                                                           {1, 2}
          });
          SmallVector<OpFoldResult> outputOffsets {batch, row, column};
          SmallVector<OpFoldResult> outputSizes = {
            rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
          Value next =
            tensor::InsertSliceOp::create(
              rewriter, nestedLoc, expanded, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
              .getResult();
          yielded.push_back(next);
          return success();
        });
      if (failed(loop))
        return failure();
      spatial::SpatYieldOp::create(rewriter, loc, loop->results.front());
      return success();
    });
  if (failed(computeOp))
    return failure();
  return computeOp->getResult(0);
}

static Value extractBatchedReductionPiece(Value partialPiecesArg,
                                          Value batch,
                                          Value hSlice,
                                          int64_t kSlice,
                                          RankedTensorType pieceType,
                                          int64_t numKSlices,
                                          int64_t numOutHSlices,
                                          int64_t numOutRows,
                                          PatternRewriter& rewriter,
                                          Location loc) {
  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
  Value batchOffset = affineMulConst(rewriter, loc, batch, numOutRows * numKSlices * numOutHSlices, anchorOp);
  Value hOffset = affineMulConst(rewriter, loc, hSlice, numKSlices * numOutRows, anchorOp);
  Value kOffset = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), kSlice * numOutRows);
  Value batchAndHSlice = arith::AddIOp::create(rewriter, loc, batchOffset, hOffset);
  Value pieceOffset = arith::AddIOp::create(rewriter, loc, batchAndHSlice, kOffset);
  SmallVector<OpFoldResult> offsets {pieceOffset, rewriter.getIndexAttr(0)};
  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(crossbarSize.getValue())};
  return tensor::ExtractSliceOp::create(
    rewriter, loc, pieceType, partialPiecesArg, offsets, sizes, getUnitStrides(rewriter, 2));
}

static Value reduceBatchedPartialPiecesForHSlice(Value partialPiecesArg,
                                                 Value batch,
                                                 Value hSlice,
                                                 RankedTensorType pieceType,
                                                 int64_t numKSlices,
                                                 int64_t numOutHSlices,
                                                 int64_t numOutRows,
                                                 PatternRewriter& rewriter,
                                                 Location loc) {
  SmallVector<Value> activePieces;
  activePieces.reserve(numKSlices);
  for (int64_t kSlice = 0; kSlice < numKSlices; ++kSlice)
    activePieces.push_back(extractBatchedReductionPiece(
      partialPiecesArg, batch, hSlice, kSlice, pieceType, numKSlices, numOutHSlices, numOutRows, rewriter, loc));

  while (activePieces.size() > 1) {
    SmallVector<Value> nextPieces;
    nextPieces.reserve((activePieces.size() + 1) / 2);
    for (size_t pieceIndex = 0; pieceIndex + 1 < activePieces.size(); pieceIndex += 2)
      nextPieces.push_back(
        spatial::SpatVAddOp::create(rewriter, loc, pieceType, activePieces[pieceIndex], activePieces[pieceIndex + 1])
          .getResult());
    if (activePieces.size() % 2 != 0)
      nextPieces.push_back(activePieces.back());
    activePieces = std::move(nextPieces);
  }

  return activePieces.front();
}

static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
                                                      RankedTensorType partialPiecesType,
                                                      RankedTensorType outType,
                                                      RankedTensorType paddedOutType,
                                                      int64_t numBatches,
                                                      int64_t numKSlices,
                                                      PatternRewriter& rewriter,
                                                      Location loc) {
  auto computeOp = createSpatCompute<1>(
    rewriter, loc, TypeRange {outType}, {}, ValueRange {partialPieces}, [&](Value partialPiecesArg) -> LogicalResult {
      const int64_t numOutRows = outType.getDimSize(1);
      const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(2), crossbarSize.getValue());
      auto pieceType = RankedTensorType::get({numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
                                             partialPiecesType.getElementType());
      auto outputSliceType = RankedTensorType::get({1, numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
                                                   partialPiecesType.getElementType());

      Value outputInit =
        tensor::EmptyOp::create(rewriter, loc, paddedOutType.getShape(), paddedOutType.getElementType()).getResult();
      Value c0 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 0);
      Value c1 = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), 1);
      Value cNumBatches = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numBatches);
      Value cNumOutHSlices =
        getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), numOutHSlices);

      auto batchLoop = buildNormalizedScfFor(
        rewriter,
        loc,
        c0,
        cNumBatches,
        c1,
        ValueRange {outputInit},
        [&](
          OpBuilder&, Location batchLoc, Value batch, ValueRange batchIterArgs, SmallVectorImpl<Value>& batchYielded) {
          auto hLoop = buildNormalizedScfFor(
            rewriter,
            batchLoc,
            c0,
            cNumOutHSlices,
            c1,
            ValueRange {batchIterArgs.front()},
            [&](OpBuilder&, Location hLoc, Value hSlice, ValueRange hIterArgs, SmallVectorImpl<Value>& hYielded) {
              Value outputAcc = hIterArgs.front();
              Value reduced = reduceBatchedPartialPiecesForHSlice(
                partialPiecesArg, batch, hSlice, pieceType, numKSlices, numOutHSlices, numOutRows, rewriter, hLoc);
              Value expandedReduced = tensor::ExpandShapeOp::create(rewriter,
                                                                    hLoc,
                                                                    outputSliceType,
                                                                    reduced,
                                                                    SmallVector<ReassociationIndices> {
                                                                      {0, 1},
                                                                      {2}
              });
              Value hOffset = affineMulConst(
                rewriter, hLoc, hSlice, crossbarSize.getValue(), rewriter.getInsertionBlock()->getParentOp());
              SmallVector<OpFoldResult> outputOffsets {batch, rewriter.getIndexAttr(0), hOffset};
              SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1),
                                                     rewriter.getIndexAttr(numOutRows),
                                                     rewriter.getIndexAttr(crossbarSize.getValue())};
              Value next =
                tensor::InsertSliceOp::create(
                  rewriter, hLoc, expandedReduced, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
                  .getResult();
              hYielded.push_back(next);
              return success();
            });
          if (failed(hLoop))
            return failure();
          batchYielded.push_back(hLoop->results.front());
          return success();
        });
      if (failed(batchLoop))
        return failure();
      Value paddedOutput = batchLoop->results.front();
      Value result = paddedOutput;
      if (paddedOutType != outType) {
        SmallVector<OpFoldResult> outputOffsets {
          rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
        SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(numBatches),
                                               rewriter.getIndexAttr(outType.getDimSize(1)),
                                               rewriter.getIndexAttr(outType.getDimSize(2))};
        result = tensor::ExtractSliceOp::create(
          rewriter, loc, outType, paddedOutput, outputOffsets, outputSizes, getUnitStrides(rewriter, 3));
      }
      spatial::SpatYieldOp::create(rewriter, loc, result);
      return success();
    });
  if (failed(computeOp))
    return failure();
  return computeOp->getResult(0);
}

struct MatMulShapeInfo {
  RankedTensorType lhsType;
  RankedTensorType rhsType;
  RankedTensorType outType;
  SmallVector<int64_t> batchShape;
  int64_t lhsBatch;
  int64_t rhsBatch;
  int64_t batch;
  int64_t m;
  int64_t k;
  int64_t n;
};

static FailureOr<MatMulShapeInfo> analyzeMatMulShape(ONNXMatMulOp matmulOp) {
  auto lhsType = dyn_cast<RankedTensorType>(matmulOp.getA().getType());
  auto rhsType = dyn_cast<RankedTensorType>(matmulOp.getB().getType());
  auto outType = dyn_cast<RankedTensorType>(matmulOp.getY().getType());
  if (!lhsType || !rhsType || !outType || !lhsType.hasStaticShape() || !rhsType.hasStaticShape()
      || !outType.hasStaticShape())
    return failure();
  if (lhsType.getRank() < 2 || rhsType.getRank() < 2 || outType.getRank() < 2)
    return failure();
  if (!hasStaticPositiveShape(lhsType) || !hasStaticPositiveShape(rhsType) || !hasStaticPositiveShape(outType))
    return failure();

  SmallVector<int64_t> lhsBatchShape(lhsType.getShape().begin(), lhsType.getShape().end() - 2);
  SmallVector<int64_t> rhsBatchShape(rhsType.getShape().begin(), rhsType.getShape().end() - 2);
  auto batchShape = inferSupportedBatchShape(lhsBatchShape, rhsBatchShape);
  if (failed(batchShape))
    return failure();

  const int64_t lhsBatch = lhsBatchShape.empty() ? 1 : getStaticShapeElementCount(lhsBatchShape);
  const int64_t rhsBatch = rhsBatchShape.empty() ? 1 : getStaticShapeElementCount(rhsBatchShape);
  const int64_t batch = batchShape->empty() ? 1 : getStaticShapeElementCount(*batchShape);
  const int64_t m = lhsType.getDimSize(lhsType.getRank() - 2);
  const int64_t k = lhsType.getDimSize(lhsType.getRank() - 1);
  const int64_t rhsK = rhsType.getDimSize(rhsType.getRank() - 2);
  const int64_t n = rhsType.getDimSize(rhsType.getRank() - 1);
  if (k != rhsK)
    return failure();

  if (outType.getRank() == 2) {
    if (batch != 1 || outType.getDimSize(0) != m || outType.getDimSize(1) != n)
      return failure();
  }
  else {
    SmallVector<int64_t> outBatchShape(outType.getShape().begin(), outType.getShape().end() - 2);
    if (!llvm::equal(outBatchShape, *batchShape) || outType.getDimSize(outType.getRank() - 2) != m
        || outType.getDimSize(outType.getRank() - 1) != n)
      return failure();
  }

  return MatMulShapeInfo {lhsType, rhsType, outType, *batchShape, lhsBatch, rhsBatch, batch, m, k, n};
}

struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(ONNXMatMulOp matmulOp, PatternRewriter& rewriter) const override {
    auto shapeInfo = analyzeMatMulShape(matmulOp);
    if (failed(shapeInfo) || shapeInfo->outType.getRank() != 2)
      return failure();

    Location loc = matmulOp.getLoc();
    bool useTransposedForm = isCompileTimeComputable(matmulOp.getA()) && !isCompileTimeComputable(matmulOp.getB());

    Value lhs = collapseBatchDims(matmulOp.getA(), shapeInfo->lhsBatch, shapeInfo->m, shapeInfo->k, rewriter, loc);
    Value rhs = collapseBatchDims(matmulOp.getB(), shapeInfo->rhsBatch, shapeInfo->k, shapeInfo->n, rewriter, loc);
    int64_t lhsBatchForGemm = shapeInfo->lhsBatch;
    int64_t rhsBatchForGemm = shapeInfo->rhsBatch;
    int64_t gemmM = shapeInfo->m;
    int64_t gemmK = shapeInfo->k;
    int64_t gemmN = shapeInfo->n;
    if (useTransposedForm) {
      lhs = transposeLastTwoDims(matmulOp.getB(), rewriter, loc);
      lhsBatchForGemm = shapeInfo->rhsBatch;
      rhs = transposeLastTwoDims(matmulOp.getA(), rewriter, loc);
      rhsBatchForGemm = shapeInfo->lhsBatch;
      gemmM = shapeInfo->n;
      gemmN = shapeInfo->m;
    }

    auto gemmType = RankedTensorType::get({gemmM, gemmN}, shapeInfo->outType.getElementType());
    Value none = ONNXNoneOp::create(rewriter, loc, rewriter.getNoneType());
    Value lhsMatrix = extractBatchMatrix(lhs, /*batchIndex=*/0, lhsBatchForGemm, gemmM, gemmK, rewriter, loc);
    Value rhsMatrix = extractBatchMatrix(rhs, /*batchIndex=*/0, rhsBatchForGemm, gemmK, gemmN, rewriter, loc);
    Value gemmResult = ONNXGemmOp::create(rewriter,
                                          loc,
                                          gemmType,
                                          lhsMatrix,
                                          rhsMatrix,
                                          none,
                                          rewriter.getF32FloatAttr(1.0f),
                                          rewriter.getF32FloatAttr(1.0f),
                                          rewriter.getBoolAttr(false),
                                          rewriter.getBoolAttr(false))
                         .getY();
    if (useTransposedForm)
      gemmResult =
        ONNXTransposeOp::create(rewriter, loc, shapeInfo->outType, gemmResult, rewriter.getI64ArrayAttr({1, 0}))
          .getResult();
    rewriter.replaceOp(matmulOp, gemmResult);
    return success();
  }
};

struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(ONNXMatMulOp matmulOp, PatternRewriter& rewriter) const override {
    auto shapeInfo = analyzeMatMulShape(matmulOp);
    if (failed(shapeInfo))
      return failure();
    if (shapeInfo->outType.getRank() == 2)
      return failure();

    Location loc = matmulOp.getLoc();
    bool useTransposedForm = isCompileTimeComputable(matmulOp.getA()) && !isCompileTimeComputable(matmulOp.getB());

    Value lhs = collapseBatchDims(matmulOp.getA(), shapeInfo->lhsBatch, shapeInfo->m, shapeInfo->k, rewriter, loc);
    Value rhs = collapseBatchDims(matmulOp.getB(), shapeInfo->rhsBatch, shapeInfo->k, shapeInfo->n, rewriter, loc);
    int64_t lhsBatchForGemm = shapeInfo->lhsBatch;
    int64_t rhsBatchForGemm = shapeInfo->rhsBatch;
    int64_t gemmM = shapeInfo->m;
    int64_t gemmK = shapeInfo->k;
    int64_t gemmN = shapeInfo->n;
    if (useTransposedForm) {
      lhs = transposeLastTwoDims(matmulOp.getB(), rewriter, loc);
      lhsBatchForGemm = shapeInfo->rhsBatch;
      rhs = transposeLastTwoDims(matmulOp.getA(), rewriter, loc);
      rhsBatchForGemm = shapeInfo->lhsBatch;
      gemmM = shapeInfo->n;
      gemmN = shapeInfo->m;
    }

    lhs = ensureBatchedTensor(lhs, lhsBatchForGemm, gemmM, gemmK, rewriter, loc);
    rhs = ensureBatchedTensor(rhs, rhsBatchForGemm, gemmK, gemmN, rewriter, loc);
    auto lhsBatchedType = cast<RankedTensorType>(lhs.getType());
    auto rhsBatchedType = cast<RankedTensorType>(rhs.getType());
    auto directOutType = RankedTensorType::get({shapeInfo->batch, gemmM, gemmN}, shapeInfo->outType.getElementType());

    if (isCompileTimeComputable(rhs)) {
      const int64_t numKSlices = ceilIntegerDivide(gemmK, crossbarSize.getValue());
      const int64_t numOutHSlices = ceilIntegerDivide(gemmN, crossbarSize.getValue());
      const int64_t paddedReductionSize = numKSlices * static_cast<int64_t>(crossbarSize.getValue());
      const int64_t paddedOutCols = numOutHSlices * static_cast<int64_t>(crossbarSize.getValue());
      auto paddedLhsType = RankedTensorType::get(
        {lhsBatchForGemm, gemmM, paddedReductionSize}, lhsBatchedType.getElementType(), lhsBatchedType.getEncoding());
      auto paddedRhsType = RankedTensorType::get({shapeInfo->batch, paddedReductionSize, paddedOutCols},
                                                 rhsBatchedType.getElementType(),
                                                 rhsBatchedType.getEncoding());
      auto paddedOutType =
        RankedTensorType::get({shapeInfo->batch, gemmM, paddedOutCols}, shapeInfo->outType.getElementType());

      auto paddedRhs = materializePaddedBatchedWeight(rhs, rhsBatchForGemm, shapeInfo->batch, paddedRhsType, rewriter);
      if (succeeded(paddedRhs)) {
        Value paddedLhs = createPaddedBatchedInputCompute(lhs, paddedLhsType, rewriter, loc);
        const int64_t laneCount = shapeInfo->batch * gemmM * numKSlices * numOutHSlices;
        auto partialPiecesType = RankedTensorType::get({laneCount, static_cast<int64_t>(crossbarSize.getValue())},
                                                       shapeInfo->outType.getElementType());
        auto batchOp = createBatchedVmmBatch(paddedLhs,
                                             *paddedRhs,
                                             paddedLhsType,
                                             lhsBatchForGemm,
                                             paddedRhsType,
                                             rhsBatchForGemm,
                                             partialPiecesType,
                                             gemmM,
                                             numKSlices,
                                             numOutHSlices,
                                             rewriter,
                                             loc);
        if (failed(batchOp))
          return failure();
        auto result = createBatchedReductionCompute(batchOp->getResult(0),
                                                    partialPiecesType,
                                                    directOutType,
                                                    paddedOutType,
                                                    shapeInfo->batch,
                                                    numKSlices,
                                                    rewriter,
                                                    loc);
        if (failed(result))
          return failure();
        Value finalResult = *result;
        if (useTransposedForm) {
          auto transposedOutType = RankedTensorType::get({shapeInfo->batch, shapeInfo->m, shapeInfo->n},
                                                         shapeInfo->outType.getElementType(),
                                                         shapeInfo->outType.getEncoding());
          finalResult =
            ONNXTransposeOp::create(rewriter, loc, transposedOutType, finalResult, rewriter.getI64ArrayAttr({0, 2, 1}))
              .getResult();
        }
        finalResult = expandBatchDims(finalResult, shapeInfo->outType, shapeInfo->batchShape.size(), rewriter, loc);
        rewriter.replaceOp(matmulOp, finalResult);
        return success();
      }
    }
    const int64_t laneCount = shapeInfo->batch * gemmM * gemmN;
    auto scalarPiecesType = RankedTensorType::get({laneCount, 1}, shapeInfo->outType.getElementType());
    auto batchOp = createBatchedVvdmulBatch(lhs,
                                            lhsBatchForGemm,
                                            rhs,
                                            rhsBatchForGemm,
                                            lhsBatchedType,
                                            rhsBatchedType,
                                            scalarPiecesType,
                                            directOutType,
                                            rewriter,
                                            loc);
    if (failed(batchOp))
      return failure();
    auto result =
      createBatchedDynamicOutputCompute(batchOp->getResult(0), scalarPiecesType, directOutType, rewriter, loc);
    if (failed(result))
      return failure();
    Value finalResult = *result;
    if (useTransposedForm) {
      auto transposedOutType = RankedTensorType::get({shapeInfo->batch, shapeInfo->m, shapeInfo->n},
                                                     shapeInfo->outType.getElementType(),
                                                     shapeInfo->outType.getEncoding());
      finalResult =
        ONNXTransposeOp::create(rewriter, loc, transposedOutType, finalResult, rewriter.getI64ArrayAttr({0, 2, 1}))
          .getResult();
    }
    finalResult = expandBatchDims(finalResult, shapeInfo->outType, shapeInfo->batchShape.size(), rewriter, loc);
    rewriter.replaceOp(matmulOp, finalResult);
    return success();
  }
};

} // namespace

void populateMatMulRewritePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
  patterns.insert<MatMulToGemm, MatMulBatchedToSpatialComputes>(ctx);
}

} // namespace onnx_mlir