//===- VectorToAMX.cpp - Convert vector to AMX dialect ----------*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "mlir/Conversion/VectorToAMX/VectorToAMX.h"

#include "mlir/Dialect/AMX/AMXDialect.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/Builders.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

#include "llvm/Support/DebugLog.h"

#include <numeric>

namespace mlir {
#define GEN_PASS_DEF_CONVERTVECTORTOAMX
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir

using namespace mlir;

#define DEBUG_TYPE "vector-to-amx"

namespace {

/// Return true if vector shape is compatible with AMX tiles.
/// The validation accounts for VNNI packing.
static bool verifyAmxShape(VectorType vec) {
  // Check overall shape:
  //   - 2D for plain layout input or output
  //   - 3D for VNNI packed input
  if (vec.getRank() != 2 && vec.getRank() != 3)
    return false;

  ArrayRef<int64_t> shape = vec.getShape();
  int64_t rows = shape[0];
  int64_t cols = shape[1];
  unsigned elemBitWidth = vec.getElementType().getIntOrFloatBitWidth();

  // 3D shape indicates VNNI packed layout.
  if (vec.getRank() == 3) {
    int64_t vnniFactor = 32 / elemBitWidth;
    if (shape.back() != vnniFactor) {
      LDBG() << "invalid VNNI packing factor";
      return false;
    }
    cols *= vnniFactor;
  }

  // AMX tile supports up to 16 rows of 64 bytes each.
  constexpr unsigned maxRows = 16;
  constexpr unsigned maxBitsPerRow = 64 * 8;
  return rows <= maxRows && (cols * elemBitWidth) <= maxBitsPerRow;
}

/// Check if contraction operands are in AMX-compatible packed VNNI layout.
static LogicalResult isAmxVnniLayout(PatternRewriter &rewriter,
                                     vector::ContractionOp contractOp) {
  VectorType accType = dyn_cast<VectorType>(contractOp.getAcc().getType());
  if (!accType || accType.getRank() != 2)
    return rewriter.notifyMatchFailure(contractOp, "Expects acc 2D vector");

  // Expect 3D inputs for VNNI packed data.
  VectorType lhsType = contractOp.getLhs().getType();
  VectorType rhsType = contractOp.getRhs().getType();
  if (lhsType.getRank() != 3 || rhsType.getRank() != 3)
    return rewriter.notifyMatchFailure(contractOp,
                                       "Expects lhs and rhs 3D vectors");

  // Check if shapes are compatible with AMX tile.
  if (!verifyAmxShape(lhsType) || !verifyAmxShape(rhsType) ||
      !verifyAmxShape(accType))
    return rewriter.notifyMatchFailure(contractOp, "Invalid operand shape");

  // Validate affine maps.
  //
  // Iterators can be ordered arbitrarily. Indexing map positions are based on
  // operands' target shapes.
  // The matrix layouts must match the following:
  //   - matrix A - [M]x[K/vnniFactor]x[vnniFactor]
  //   - matrix B - [K/vnniFactor]x[N]x[vnniFactor]
  //   - matrix C - [M]x[N]
  SmallVector<AffineMap, 4> indexingMaps = contractOp.getIndexingMapsArray();
  AffineMap mapA = indexingMaps[0];
  AffineMap mapB = indexingMaps[1];
  if (mapA.getNumInputs() != 4 || mapA.getNumResults() != 3 ||
      mapB.getNumResults() != 3)
    return rewriter.notifyMatchFailure(contractOp,
                                       "Invalid input indexing maps");
  FailureOr<linalg::ContractionDimensions> dims =
      linalg::inferContractionDims(indexingMaps);
  if (failed(dims))
    return rewriter.notifyMatchFailure(contractOp,
                                       "Failed to infer contraction dims");
  // Two reduction dimensions are expected:
  //   - one for the K dimension
  //   - one for the VNNI factor
  if (dims->k.size() != 2)
    return rewriter.notifyMatchFailure(contractOp,
                                       "Expected two reduction dims");
  assert(dims->m.size() == 1 && dims->n.size() == 1 &&
         "Invalid parallel contraction dims");

  SmallVector<vector::IteratorType> iteratorTypes =
      contractOp.getIteratorTypesArray();
  // Check VNNI dim maps - the innermost dim for A and B inputs.
  auto vnniDimA = dyn_cast<AffineDimExpr>(mapA.getResult(2));
  auto vnniDimB = dyn_cast<AffineDimExpr>(mapB.getResult(2));
  if (!vnniDimA || !vnniDimB || vnniDimA != vnniDimB ||
      iteratorTypes[vnniDimA.getPosition()] != vector::IteratorType::reduction)
    return rewriter.notifyMatchFailure(contractOp, "Invalid VNNI dim map");
  // Check K dim maps - non-transposed row-major layout.
  auto redDimA = dyn_cast<AffineDimExpr>(mapA.getResult(1));
  auto redDimB = dyn_cast<AffineDimExpr>(mapB.getResult(0));
  if (!redDimA || !redDimB || redDimA != redDimB ||
      iteratorTypes[redDimA.getPosition()] != vector::IteratorType::reduction)
    return rewriter.notifyMatchFailure(contractOp, "Invalid K dim map");
  // Check M and N dim maps - map to non-transposed output.
  AffineMap mapC = indexingMaps[2];
  auto mDimC = dyn_cast<AffineDimExpr>(mapC.getResult(0));
  auto nDimC = dyn_cast<AffineDimExpr>(mapC.getResult(1));
  if (!mDimC || !nDimC)
    return rewriter.notifyMatchFailure(contractOp, "Invalid acc maps");
  auto parallelDimA = dyn_cast<AffineDimExpr>(mapA.getResult(0));
  if (!parallelDimA ||
      iteratorTypes[parallelDimA.getPosition()] !=
          vector::IteratorType::parallel ||
      parallelDimA != mDimC)
    return rewriter.notifyMatchFailure(contractOp, "Invalid M dim map");
  auto parallelDimB = dyn_cast<AffineDimExpr>(mapB.getResult(1));
  if (!parallelDimB ||
      iteratorTypes[parallelDimB.getPosition()] !=
          vector::IteratorType::parallel ||
      parallelDimB != nDimC)
    return rewriter.notifyMatchFailure(contractOp, "Invalid N dim map");

  return success();
}

/// Validate contraction operands for AMX lowering.
static LogicalResult validateOperands(PatternRewriter &rewriter,
                                      vector::ContractionOp contractOp) {
  VectorType accType = dyn_cast<VectorType>(contractOp.getAcc().getType());
  if (!accType)
    return rewriter.notifyMatchFailure(contractOp, "Expects vector acc");

  // Check if operand types are compatible with AMX compute ops.
  bool validElemTypes = false;
  Type lhsElemType = contractOp.getLhs().getType().getElementType();
  Type rhsElemType = contractOp.getRhs().getType().getElementType();
  Type accElemType = accType.getElementType();
  if (accElemType.isInteger(32)) {
    validElemTypes = lhsElemType.isInteger(8) && rhsElemType.isInteger(8);
  } else if (accElemType.isF32()) {
    validElemTypes = (lhsElemType.isF16() && rhsElemType.isF16()) ||
                     (lhsElemType.isBF16() && rhsElemType.isBF16());
  }
  if (!validElemTypes)
    return rewriter.notifyMatchFailure(contractOp,
                                       "Invalid combination of operand types");

  if (failed(isAmxVnniLayout(rewriter, contractOp)))
    return failure();

  return success();
}

/// Collapse the two innermost dimensions together.
static TypedValue<MemRefType> collapseLastDim(PatternRewriter &rewriter,
                                              TypedValue<MemRefType> memref) {
  int64_t rank = memref.getType().getRank();
  SmallVector<ReassociationIndices> reassocIndices;
  for (auto i : llvm::seq<int64_t>(0, rank - 2))
    reassocIndices.push_back({i});
  reassocIndices.push_back({rank - 2, rank - 1});
  return memref::CollapseShapeOp::create(rewriter, memref.getLoc(), memref,
                                         reassocIndices);
}

/// Attempt to create an AMX tile load/store operation equivalent to the given
/// vector transfer `xfer` op.
/// This approach allows to skip longer route through registers and a temporary
/// buffer otherwise required to move data to/from an AMX tile.
static Operation *
loadStoreFromTransfer(PatternRewriter &rewriter,
                      VectorTransferOpInterface xferOp, bool isPacked,
                      TypedValue<amx::TileType> tileToStore = nullptr) {
  if (!xferOp || !isa<vector::TransferReadOp, vector::TransferWriteOp>(xferOp))
    return nullptr;
  if (xferOp.hasOutOfBoundsDim() ||
      !xferOp.getPermutationMap().isMinorIdentity())
    return nullptr;

  // Extra checks in case of a write op.
  // Stores must not be packed.
  if (isa<vector::TransferWriteOp>(xferOp) &&
      (!tileToStore || isPacked ||
       tileToStore.getType().getShape() != xferOp.getVectorType().getShape()))
    return nullptr;

  // Check for a memref source buffer.
  // AMX data transfer requires at least 2D shape to correctly
  // infer stride between rows.
  Value base = xferOp.getBase();
  auto memTy = dyn_cast<MemRefType>(base.getType());
  int64_t memRank = memTy.getRank();
  if (!memTy || memRank < 2)
    return nullptr;

  // Check that the source buffer has enough contiguous elements to load whole
  // AMX tile row.
  //
  // To ensure correctness, the validation is conservative and expects the
  // buffer's innermost dimensions to be statically known, equal to or larger
  // than the vector row length, and equal to the VNNI dimension if applicable.
  //
  // This check could be relaxed to accept more arbitrarily shaped buffers as
  // long as there are enough contiguous elements to load a whole row.
  if (!memTy.areTrailingDimsContiguous(isPacked ? 2 : 1))
    return nullptr;
  VectorType vecTy = xferOp.getVectorType();
  ArrayRef<int64_t> vecShape = vecTy.getShape();
  ArrayRef<int64_t> memShape = memTy.getShape();
  if (memShape.back() == ShapedType::kDynamic ||
      memShape.back() < vecShape.back())
    return nullptr;
  if (isPacked &&
      (memShape.back() != vecShape.back() ||
       memShape[memShape.size() - 2] == ShapedType::kDynamic ||
       memShape[memShape.size() - 2] < vecShape[vecShape.size() - 2]))
    return nullptr;

  // Load values directly from the buffer to an AMX tile.
  PatternRewriter::InsertionGuard g(rewriter);
  rewriter.setInsertionPoint(xferOp);
  Location loc = xferOp.getLoc();

  // Create a subview of the source buffer based on the transfer op to resolve
  // offsets.
  SmallVector<OpFoldResult> strides(memRank, rewriter.getIndexAttr(1));
  int64_t vecRank = vecTy.getRank();
  assert(memRank >= vecRank &&
         "Expects buffer to be the same or greater rank than vector");
  SmallVector<int64_t> shape(memRank - vecRank, 1);
  shape.append(vecShape.begin(), vecShape.end());
  TypedValue<MemRefType> src =
      memref::SubViewOp::create(
          rewriter, loc, base, getAsOpFoldResult(xferOp.getIndices()),
          getAsOpFoldResult(rewriter.getI64ArrayAttr(shape)), strides)
          .getResult();

  // Collapse the VNNI dimension in case of packing.
  if (isPacked)
    src = collapseLastDim(rewriter, src);
  int64_t rows = vecShape[0];
  int64_t cols = std::accumulate(vecShape.begin() + 1, vecShape.end(), 1,
                                 std::multiplies<int64_t>());
  auto tileType = amx::TileType::get({rows, cols}, vecTy.getElementType());

  Value zeroIndex = rewriter.createOrFold<arith::ConstantIndexOp>(loc, 0);
  SmallVector<Value> tileIndicides(src.getType().getRank(), zeroIndex);

  Operation *amxTileOp = nullptr;
  if (isa<vector::TransferReadOp>(xferOp)) {
    amxTileOp =
        amx::TileLoadOp::create(rewriter, loc, tileType, src, tileIndicides);
  } else if (isa<vector::TransferWriteOp>(xferOp)) {
    amxTileOp = amx::TileStoreOp::create(rewriter, loc, src, tileIndicides,
                                         tileToStore);
  } else {
    llvm_unreachable("unsupported vector transfer op");
  }

  return amxTileOp;
}

/// Attempt to create an AMX tile load operation equivalent to the given
/// vector transfer `readOp`.
/// Returns loaded AMX tile if successful.
static FailureOr<TypedValue<amx::TileType>>
loadFromTransfer(PatternRewriter &rewriter, vector::TransferReadOp readOp,
                 bool isPacked) {
  amx::TileLoadOp loadOp = dyn_cast_if_present<amx::TileLoadOp>(
      loadStoreFromTransfer(rewriter, readOp, isPacked));
  if (!loadOp)
    return failure();
  return loadOp.getRes();
}

/// Attempt to create an AMX tile store operation equivalent to the given
/// vector transfer `writeOp`.
static LogicalResult storeFromTransfer(PatternRewriter &rewriter,
                                       vector::TransferWriteOp writeOp,
                                       TypedValue<amx::TileType> tileToStore) {
  return success(loadStoreFromTransfer(rewriter, writeOp, /*isPacked=*/false,
                                       tileToStore));
}

/// Load vector values to an AMX tile.
static TypedValue<amx::TileType> loadTile(PatternRewriter &rewriter,
                                          TypedValue<VectorType> vec) {
  Location loc = vec.getLoc();

  VectorType vecTy = vec.getType();
  bool isPacked = vecTy.getRank() == 3;

  // Try to load tile directly from vector producer's buffer.
  auto readOp = vec.getDefiningOp<vector::TransferReadOp>();
  FailureOr<TypedValue<amx::TileType>> tile =
      loadFromTransfer(rewriter, readOp, isPacked);
  if (succeeded(tile))
    return *tile;

  // Transfer the vector to a tile through an intermediate buffer.
  Value buf = memref::AllocaOp::create(
      rewriter, loc, MemRefType::get(vecTy.getShape(), vecTy.getElementType()));
  Value zeroIndex = rewriter.createOrFold<arith::ConstantIndexOp>(loc, 0);
  SmallVector<Value> indices(vecTy.getRank(), zeroIndex);
  vector::TransferWriteOp::create(rewriter, loc, vec, buf, indices);

  // Collapse the VNNI dimension in case of packing.
  if (isPacked)
    buf = collapseLastDim(rewriter, cast<TypedValue<MemRefType>>(buf));

  ArrayRef<int64_t> shape = vecTy.getShape();
  int64_t rows = shape[0];
  int64_t cols = std::accumulate(shape.begin() + 1, shape.end(), 1,
                                 std::multiplies<int64_t>());
  auto tileType = amx::TileType::get({rows, cols}, vecTy.getElementType());

  return amx::TileLoadOp::create(rewriter, loc, tileType, buf,
                                 {zeroIndex, zeroIndex});
}

/// Store an AMX tile in a vector.
static TypedValue<VectorType> storeTile(PatternRewriter &rewriter,
                                        TypedValue<amx::TileType> tile) {
  Location loc = tile.getLoc();

  // Transfer the tile to a vector through an intermediate buffer.
  amx::TileType tileTy = tile.getType();
  Value buf = memref::AllocaOp::create(
      rewriter, loc,
      MemRefType::get(tileTy.getShape(), tileTy.getElementType()));
  Value zeroIndex = rewriter.createOrFold<arith::ConstantIndexOp>(loc, 0);
  SmallVector<Value> indices(2, zeroIndex);
  amx::TileStoreOp::create(rewriter, loc, buf, indices, tile);

  auto vecTy = VectorType::get(tileTy.getShape(), tileTy.getElementType());
  return vector::TransferReadOp::create(rewriter, loc, vecTy, buf, indices, {});
}

struct ContractionToAMX : public OpRewritePattern<vector::ContractionOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
                                PatternRewriter &rewriter) const override {
    Location loc = contractOp.getLoc();

    if (contractOp.getKind() != vector::CombiningKind::ADD)
      return rewriter.notifyMatchFailure(contractOp,
                                         "Expects add combining kind");
    if (failed(validateOperands(rewriter, contractOp)))
      return failure();

    TypedValue<amx::TileType> lhsTile = loadTile(rewriter, contractOp.getLhs());
    TypedValue<amx::TileType> rhsTile = loadTile(rewriter, contractOp.getRhs());
    auto acc = dyn_cast<TypedValue<VectorType>>(contractOp.getAcc());
    assert(acc && "Invalid accumulator type");
    TypedValue<amx::TileType> accTile = loadTile(rewriter, acc);

    TypedValue<amx::TileType> tileMul;
    if (acc.getType().getElementType().isFloat()) {
      tileMul = amx::TileMulFOp::create(rewriter, loc, accTile.getType(),
                                        lhsTile, rhsTile, accTile);
    } else {
      tileMul = amx::TileMulIOp::create(rewriter, loc, accTile.getType(),
                                        lhsTile, rhsTile, accTile);
    }

    // If the contraction result is only written back to memory, try to replace
    // the vector op with an AMX store directly.
    Value res = contractOp.getResult();
    if (res.hasOneUse()) {
      auto writeOp = dyn_cast<vector::TransferWriteOp>(*res.getUsers().begin());
      LogicalResult storeRes = storeFromTransfer(rewriter, writeOp, tileMul);
      if (succeeded(storeRes)) {
        rewriter.eraseOp(writeOp);
        rewriter.eraseOp(contractOp);
        return success();
      }
    }

    // Load the result back into a vector.
    Value newResult = storeTile(rewriter, tileMul);
    rewriter.replaceOp(contractOp, newResult);

    return success();
  }
};

struct ConvertVectorToAMXPass
    : public impl::ConvertVectorToAMXBase<ConvertVectorToAMXPass> {
  void runOnOperation() override {
    MLIRContext &ctx = getContext();
    RewritePatternSet patterns(&ctx);
    populateVectorToAMXConversionPatterns(patterns);
    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
      return signalPassFailure();
  }
};

} // namespace

void mlir::populateVectorToAMXConversionPatterns(RewritePatternSet &patterns) {
  patterns.add<ContractionToAMX>(patterns.getContext());
}