// Copyright 2020 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include "iree/compiler/Codegen/PassDetail.h"
#include "iree/compiler/Codegen/Passes.h"
#include "iree/compiler/Utils/CustomKernelsTargetInfo.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/Triple.h"
#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
#include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
#include "mlir/Dialect/ArmNeon/ArmNeonDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/Builders.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

namespace mlir {
namespace iree_compiler {

namespace {

// Returns true if `contractionOp` is of the form
//   matrix * transposed_matrix.
// That is, if there are 2 parallel iterators, say M and N, 1 additive reduction
// iterator, say K, and the indexing maps are {{M, K}, {N, K}, {M, N}}.
static bool isMatrixTimesMatrixTransposed(vector::ContractionOp contractionOp) {
  // Check that the reduction is additive.
  if (contractionOp.kind() != vector::CombiningKind::ADD) {
    return false;
  }
  // Check that there are 2 parallel and 1 reduction iterators.
  auto iteratorTypes = contractionOp.iterator_types().getValue();
  if (iteratorTypes.size() != 3) {
    return false;
  }
  SmallVector<int, 3> parallel_iterators;
  SmallVector<int, 3> reduction_iterators;
  for (int i = 0; i < 3; i++) {
    if (isParallelIterator(iteratorTypes[i])) {
      parallel_iterators.push_back(i);
    } else if (isReductionIterator(iteratorTypes[i])) {
      reduction_iterators.push_back(i);
    } else {
      return false;
    }
  }
  if (parallel_iterators.size() != 2 || reduction_iterators.size() != 1) {
    return false;
  }
  // Give the found iterators some idiomatic names.
  const int MIter = parallel_iterators[0];
  const int NIter = parallel_iterators[1];
  const int KIter = reduction_iterators[0];
  // Check that there are 3 indexing maps.
  auto indexingMaps = contractionOp.indexing_maps().getValue();
  if (indexingMaps.size() != 3) {
    return false;
  }
  // Check that the indexing maps have the expected form.
  const int expectedMapResults[3][2] = {
      {MIter, KIter}, {NIter, KIter}, {MIter, NIter}};
  for (int m = 0; m < 3; ++m) {
    auto map = indexingMaps[m].cast<AffineMapAttr>().getValue();
    if (map.getNumDims() != 3 || map.getNumResults() != 2) {
      return false;
    }
    for (int r = 0; r < 2; ++r) {
      int actualMapResult =
          map.getResults()[r].cast<AffineDimExpr>().getPosition();
      if (actualMapResult != expectedMapResults[m][r]) {
        return false;
      }
    }
  }
  return true;
}

// Returns true if `contractionOp` is of the form
//   matrix * transposed_matrix
// where matrix is a vector<{mSize}x{kSize}xType>, and
// transposed_matrix is a vector<{nSize}x{kSize}xType>
static bool isMatrixTimesMatrixTransposedOfGivenShape(
    vector::ContractionOp contractionOp, int64_t mSize, int64_t kSize,
    int64_t nSize) {
  if (!isMatrixTimesMatrixTransposed(contractionOp)) {
    return false;
  }
  VectorType lhsType = contractionOp.lhs().getType().cast<VectorType>();
  VectorType rhsType = contractionOp.rhs().getType().cast<VectorType>();
  auto lhsShape = lhsType.getShape();
  auto rhsShape = rhsType.getShape();
  if (lhsShape[0] != mSize || lhsShape[1] != kSize || rhsShape[0] != nSize ||
      rhsShape[1] != kSize) {
    return false;
  }
  return true;
}

// Checks that the Value `extResult` is defined by an arith::ExtSIOp promoting
// from `extSrcType` to `extDstType`, and returns the input of the ExtSIOp.
// Note that this only looks at the immediately defining operation, so we likely
// want to have earlier passes that sink widening operations as far down as
// possible, which is probably just good regardless.
static Value getExtSIInput(Type extSrcType, Type extDstType, Value extResult) {
  auto extSIOp = extResult.getDefiningOp<arith::ExtSIOp>();
  if (!extSIOp) {
    return nullptr;
  }
  Value extInput = extSIOp.getIn();
  if (extInput.getType().cast<VectorType>().getElementType() != extSrcType) {
    return nullptr;
  }
  return extInput;
}

// Helper to create a 1D, contiguous slice of a 1D vector.
static Value extract1DSlice(PatternRewriter &rewriter, Location loc,
                            VectorType dstVecType, Value input, int position) {
  assert(input.getType().cast<VectorType>().getRank() == 1);
  assert(dstVecType.getRank() == 1);
  std::array<int64_t, 1> offsets{position};
  std::array<int64_t, 1> strides{1};
  return rewriter.create<vector::ExtractStridedSliceOp>(
      loc, input, offsets, dstVecType.getShape(), strides);
}

// Helper to flatten a N-dimensional vector to a 1D vector.
static Value flatten(PatternRewriter &rewriter, Location loc, Value vector) {
  VectorType inputVecType = vector.getType().cast<VectorType>();
  VectorType dstType = VectorType::get(inputVecType.getNumElements(),
                                       inputVecType.getElementType());
  return rewriter.create<vector::ShapeCastOp>(loc, dstType, vector);
}

/// Converts matrix-times-matrix-transposed vector.contracts with
/// lhs and rhs inputs defined by arith.extsi promoting from i8 to i32,
///
///     %lhs_i32 = arith.extsi %lhs_i8 : i8 to i32
///     %rhs_i32 = arith.extsi %rhs_i8 : i8 to i32
///     %result = vector.contract [...]
///                 %lhs_i32 : vector<8x4xi32>,
///                 %rhs_i32 : vector<8x4xi32>,
///                 %acc_i32 : vector<8x8xi32>,
///                 [...]
///
/// To vector ops reading directly from the %lhs_i8 and %rhs_i8 values
/// (bypassing the existing arith.extsi) and passing that to a llvm.inline_asm
/// block implementing the matrix multiplication arithmetic using Aarch64
/// dot-product instructions (sdot).
struct MMT_8x4x8_i8i8i32_Aarch64Dotprod_InlineAsm
    : OpRewritePattern<vector::ContractionOp> {
  using OpRewritePattern<vector::ContractionOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(vector::ContractionOp contractionOp,
                                PatternRewriter &rewriter) const override {
    // Check if `contractionOp` matches, and obtain the un-promoted i8 input
    // LHS and RHS vectors, `lhsI8` and `rhsI8`.
    if (!isMatrixTimesMatrixTransposedOfGivenShape(contractionOp, 8, 4, 8)) {
      return failure();
    }
    Type I8Type = rewriter.getIntegerType(8);
    Type I32Type = rewriter.getIntegerType(32);
    VectorType accType = contractionOp.acc().getType().cast<VectorType>();
    if (accType.getElementType() != I32Type) {
      return failure();
    }
    Value lhsI8 = getExtSIInput(I8Type, I32Type, contractionOp.lhs());
    Value rhsI8 = getExtSIInput(I8Type, I32Type, contractionOp.rhs());
    if (!lhsI8 || !rhsI8) {
      return failure();
    }

    // `contractionOp` matches, start rewriting it. We only reference
    // the `lhsI8` and `rhsI8` values obtained above as the inputs of the
    // arith.extsi, so this rewrite will leave the existing arith.extsi without
    // any user (unless something else was using them), so they may be
    // removed by another transformation.
    Location loc = contractionOp.getLoc();
    // Flatten the inputs to 1D vectors.
    Value flatLhsI8 = flatten(rewriter, loc, lhsI8);
    Value flatRhsI8 = flatten(rewriter, loc, rhsI8);
    Value flatAcc = flatten(rewriter, loc, contractionOp.acc());

    // Create the 1D input vectors of 16 bytes each that are directly what
    // the target SIMD instructions will want.
    SmallVector<Value> lhsVec;
    SmallVector<Value> rhsVec;
    VectorType vector16xi8Type = VectorType::get({16}, I8Type);
    for (int position = 0; position < 8 * 4; position += 16) {
      lhsVec.push_back(
          extract1DSlice(rewriter, loc, vector16xi8Type, flatLhsI8, position));
      rhsVec.push_back(
          extract1DSlice(rewriter, loc, vector16xi8Type, flatRhsI8, position));
    }
    SmallVector<Value> accVec;
    VectorType int32x4Type = VectorType::get({4}, I32Type);
    for (int position = 0; position < 8 * 8; position += 4) {
      accVec.push_back(
          extract1DSlice(rewriter, loc, int32x4Type, flatAcc, position));
    }

    // Create the inline asm op's operands list.
    SmallVector<Value> asmOperands;
    // First the inputs operands.
    asmOperands.append(lhsVec);
    asmOperands.append(rhsVec);
    // Then the input-output operands.
    asmOperands.append(accVec);
    SmallVector<Type> asmOutputOperandTypes(
        llvm::map_range(accVec, [](Value v) { return v.getType(); }));

    // Create the inline asm op.
    auto returnType = LLVM::LLVMStructType::getLiteral(rewriter.getContext(),
                                                       asmOutputOperandTypes);
    auto dialectAttr = LLVM::AsmDialectAttr::get(rewriter.getContext(),
                                                 LLVM::AsmDialect::AD_ATT);
    // The LLVM inline asm syntax is documented here:
    // https://llvm.org/docs/LangRef.html#inline-assembler-expressions
    LLVM::InlineAsmOp asmOp = rewriter.create<LLVM::InlineAsmOp>(
        loc, returnType, asmOperands,
        R"ASM(
            sdot $0.4s, $18.16b, $16.4b[0]
            sdot $1.4s, $19.16b, $16.4b[0]
            sdot $2.4s, $18.16b, $16.4b[1]
            sdot $3.4s, $19.16b, $16.4b[1]
            sdot $4.4s, $18.16b, $16.4b[2]
            sdot $5.4s, $19.16b, $16.4b[2]
            sdot $6.4s, $18.16b, $16.4b[3]
            sdot $7.4s, $19.16b, $16.4b[3]
            sdot $8.4s, $18.16b, $17.4b[0]
            sdot $9.4s, $19.16b, $17.4b[0]
            sdot $10.4s, $18.16b, $17.4b[1]
            sdot $11.4s, $19.16b, $17.4b[1]
            sdot $12.4s, $18.16b, $17.4b[2]
            sdot $13.4s, $19.16b, $17.4b[2]
            sdot $14.4s, $18.16b, $17.4b[3]
            sdot $15.4s, $19.16b, $17.4b[3]
          )ASM",
        "=w,=w,=w,=w,=w,=w,=w,=w,=w,=w,=w,=w,=w,=w,=w,=w,w,w,w,w,0,1,2,3,4,5,6,"
        "7,8,9,10,11,12,13,14,15",
        /*has_side_effects=*/false, /*is_align_stack=*/false, dialectAttr,
        /*operand_attrs=*/ArrayAttr());

    // Extract result vectors from the asm op.
    SmallVector<Value, 16> resVec;
    for (int i = 0; i < 16; ++i) {
      resVec.push_back(rewriter.create<LLVM::ExtractValueOp>(
          loc, int32x4Type, asmOp.getRes(), rewriter.getI64ArrayAttr({i})));
    }

    // Insert the result vectors of size 4 into the overall result vector of
    // size 64, still 1D.
    VectorType int32x64xType = VectorType::get({64}, I32Type);
    Value result = rewriter.create<arith::ConstantOp>(
        loc, int32x64xType, DenseIntElementsAttr::get(int32x64xType, 0));
    for (int i = 0; i < 16; ++i) {
      result = rewriter.create<vector::InsertStridedSliceOp>(
          loc, resVec[i], result, std::array<int64_t, 1>{4 * i},
          std::array<int64_t, 1>{1});
    }

    // Cast the result from 1D to 2D and replace the original vector.contract.
    VectorType int32x8x8xType = VectorType::get({8, 8}, I32Type);
    rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(contractionOp,
                                                     int32x8x8xType, result);
    return success();
  }
};

/// Converts matrix-times-matrix-transposed vector.contracts with
/// lhs and rhs inputs defined by arith.extsi promoting from i8 to i32,
///
///     %lhs_i32 = arith.extsi %lhs_i8 : i8 to i32
///     %rhs_i32 = arith.extsi %rhs_i8 : i8 to i32
///     %result = vector.contract [...]
///                 %lhs_i32 : vector<8x4xi32>,
///                 %rhs_i32 : vector<8x4xi32>,
///                 %acc_i32 : vector<8x8xi32>,
///                 [...]
///
/// To vector ops reading directly from the %lhs_i8 and %rhs_i8 values
/// (bypassing the existing arith.extsi) and passing that to a llvm.inline_asm
/// block implementing the matrix multiplication arithmetic using Aarch64
/// dot-product instructions (sdot).
/// It matches the same patterns as MMT_8x4x8_i8i8i32_Aarch64Dotprod_InlineAsm
struct MMT_8x4x8_i8i8i32_Aarch64Dotprod_Intrinsics
    : public OpRewritePattern<vector::ContractionOp> {
 public:
  using OpRewritePattern<vector::ContractionOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(vector::ContractionOp contractionOp,
                                PatternRewriter &rewriter) const override {
    if (!isMatrixTimesMatrixTransposedOfGivenShape(contractionOp, 8, 4, 8)) {
      return failure();
    }

    Type I8Type = rewriter.getIntegerType(8);
    Type I32Type = rewriter.getIntegerType(32);

    auto acc = contractionOp.acc();
    auto lhs = contractionOp.lhs();
    auto rhs = contractionOp.rhs();
    if (acc.getType().cast<VectorType>().getElementType() != I32Type) {
      return failure();
    }

    Value inLhs = getExtSIInput(I8Type, I32Type, lhs);
    Value inRhs = getExtSIInput(I8Type, I32Type, rhs);

    if (!inLhs || !inRhs) return failure();

    auto loc = contractionOp.getLoc();

    auto int32x4VType = VectorType::get({4}, I32Type);

    std::array<Value, 16> accChunks;
    {
      int idx = 0;
      for (int row = 0; row < 8; ++row) {
        auto accRow = rewriter.create<vector::ExtractOp>(
            loc, acc, ArrayRef<int64_t>{row});
        for (int col = 0; col < 8; col += 4) {
          auto accChunk = rewriter.create<vector::ExtractStridedSliceOp>(
              loc, accRow, ArrayRef<int64_t>{col}, ArrayRef<int64_t>{4},
              ArrayRef<int64_t>{1});
          assert(accChunk.getType() == int32x4VType);
          accChunks[idx++] = accChunk;
        }
      }
    }

    auto int8x4x4VType = VectorType::get({4, 4}, rewriter.getIntegerType(8));
    auto extract4x4 = [&](Value in, int rowOffset, int colOffset) {
      auto chunk = rewriter.create<vector::ExtractStridedSliceOp>(
          loc, in, ArrayRef<int64_t>{rowOffset, colOffset},
          ArrayRef<int64_t>{4, 4}, ArrayRef<int64_t>{1, 1});
      assert(chunk.getType() == int8x4x4VType);
      return chunk;
    };

    std::array<Value, 2> lhsHalves = {extract4x4(inLhs, 0, 0),
                                      extract4x4(inLhs, 4, 0)};
    std::array<Value, 2> rhsHalves = {extract4x4(inRhs, 0, 0),
                                      extract4x4(inRhs, 4, 0)};

    auto int8Zero4x4 = rewriter.create<arith::ConstantOp>(
        loc, rewriter.getZeroAttr(int8x4x4VType));
    auto sdot = [&](Value acc, Value a, Value b, int64_t lane) -> Value {
      auto bReplicatedLane = rewriter.create<vector::ShuffleOp>(
          loc, b, int8Zero4x4, ArrayRef<int64_t>{lane, lane, lane, lane});

      return rewriter.create<arm_neon::Sdot2dOp>(loc, int32x4VType, acc, a,
                                                 bReplicatedLane);
    };

    std::array<Value, 16> dstChunks;
    {
      int idx = 0;
      for (Value lhs : lhsHalves) {
        for (int lane = 0; lane < 4; ++lane) {
          for (Value rhs : rhsHalves) {
            dstChunks[idx] = sdot(accChunks[idx], rhs, lhs, lane);
            ++idx;
          }
        }
      }
    }

    // Put the results back in the accumulator
    {
      int idx = 0;
      for (int row = 0; row < 8; ++row) {
        for (int col = 0; col < 8; col += 4) {
          acc = rewriter.create<vector::InsertStridedSliceOp>(
              loc, dstChunks[idx++], acc, ArrayRef<int64_t>{row, col},
              ArrayRef<int64_t>{1});
        }
      }
    }
    rewriter.replaceOp(contractionOp, {acc});
    return success();
  }
};

class VectorContractCustomKernelsPass
    : public VectorContractCustomKernelsBase<VectorContractCustomKernelsPass> {
 public:
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<vector::VectorDialect, LLVM::LLVMDialect>();
    if (target_info.has(CustomKernelTargetFeature::Intrinsics)) {
      registry.insert<arm_neon::ArmNeonDialect>();
    }
  }
  LogicalResult initializeOptions(StringRef options) override {
    if (failed(Pass::initializeOptions(options))) {
      return failure();
    }
    if (failed(ParseCustomKernelsTargetInfo(arch, features, target_info))) {
      return failure();
    }
    if (intrinsics) {
      target_info.add(CustomKernelTargetFeature::Intrinsics);
    }
    return success();
  }
  void runOnOperation() override {
    MLIRContext *context = &getContext();
    RewritePatternSet patterns(context);
    populateVectorContractCustomKernelsPatterns(target_info, patterns);
    if (failed(applyPatternsAndFoldGreedily(getOperation(),
                                            std::move(patterns)))) {
      signalPassFailure();
    }
  }

 private:
  CustomKernelsTargetInfo target_info;
};

}  // namespace

void populateVectorContractCustomKernelsPatterns(
    const CustomKernelsTargetInfo &target_info, RewritePatternSet &patterns) {
  MLIRContext *context = patterns.getContext();
  if (target_info.has(CustomKernelTargetFeature::Aarch64Dotprod)) {
    if (target_info.has(CustomKernelTargetFeature::Intrinsics)) {
      patterns.insert<MMT_8x4x8_i8i8i32_Aarch64Dotprod_Intrinsics>(context);
    } else {
      patterns.insert<MMT_8x4x8_i8i8i32_Aarch64Dotprod_InlineAsm>(context);
    }
  }
}

std::unique_ptr<OperationPass<FuncOp>> createVectorContractCustomKernelsPass() {
  return std::make_unique<VectorContractCustomKernelsPass>();
}

}  // namespace iree_compiler
}  // namespace mlir