// Copyright 2023 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/Dialect/Flow/IR/FlowOps.h"
#include "iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h"
#include "iree/compiler/GlobalOptimization/Passes.h"
#include "llvm/ADT/STLExtras.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/Tensor/Utils/Utils.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

#define DEBUG_TYPE "iree-global-opt-fuse-dequantization-matmul"
#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")

namespace mlir::iree_compiler::GlobalOptimization {

#define GEN_PASS_DEF_FUSEDEQUANTIZATIONMATMULPASS
#include "iree/compiler/GlobalOptimization/Passes.h.inc"

namespace {

//----------------------------------------------------------------------------//
//                                Utility
//----------------------------------------------------------------------------//

// Checks if the passed op is a contraction with two reduction dimensions
// This function checks that the genericOp:
// 1. isaContractionOpInterface
// 2. Has 2 reduction dimensions
static LogicalResult
isContractionWithTwoReductions(linalg::GenericOp genericOp) {
  unsigned numLoops = genericOp.getNumLoops();
  linalg::LinalgOp linalgOp = cast<linalg::LinalgOp>(genericOp.getOperation());
  if (numLoops == 0) {
    return failure();
  }
  if (!linalg::isaContractionOpInterface(linalgOp)) {
    return failure();
  }
  if (genericOp.getNumReductionLoops() != 2) {
    return failure();
  }
  return success();
}

// Checks if the passed op is a dequantization on grouped input
// This function checks that the genericOp:
// 1. Has a body like:
//      arith.extui
//      arith.uitofp
//      arith.subf
//      arith.mulf
// 2. Increases the bit width of the input
// 3. Has 3 parallel dims
// 4. Has 2 (weights, scales) or 3 (weights, scales, zero points)
//    inputs and 1 output
static LogicalResult isGroupedDequantizationOp(linalg::GenericOp genericOp) {
  // Check for 1 result, and 2 (input, scales) or 3 (input, scales, zero points)
  // inputs
  if (genericOp.getNumDpsInits() != 1) {
    return failure();
  }
  if (genericOp.getNumDpsInputs() != 2 && genericOp.getNumDpsInputs() != 3) {
    return failure();
  }
  // Check that the rank is at least 3 and all loops are parallel
  unsigned numLoops = genericOp.getNumLoops();
  unsigned numParallelLoops = genericOp.getNumParallelLoops();
  if (numLoops < 3) {
    return failure();
  }
  if (numLoops != numParallelLoops) {
    return failure();
  }

  // Work back from linalg.yield and check body of genericOp.
  // The genericOp should yield the result of an arith.mulf,
  // preceded by an arith.subf, arith.uitofp, and arith.extui
  auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
  Value producerOutput;
  Operation *producer;

  // Producer of linalg.yield op is arith.mulf
  {
    producerOutput = yieldOp->getOperand(0);
    producer = producerOutput.getDefiningOp();
    if (!producer || producer->getNumOperands() == 0) {
      return failure();
    }
    if (!matchPattern(producer, m_Op<arith::MulFOp>())) {
      return failure();
    }
  }

  // Producer of arith.mulf op is arith.subf
  {
    producerOutput = producer->getOperand(0);
    producer = producerOutput.getDefiningOp();
    if (!producer || producer->getNumOperands() == 0) {
      return failure();
    }
    if (!matchPattern(producer, m_Op<arith::SubFOp>())) {
      return failure();
    }
  }

  // Producer of arith.subf op is arith.uitofp
  {
    producerOutput = producer->getOperand(0);
    producer = producerOutput.getDefiningOp();
    if (!producer || producer->getNumOperands() == 0) {
      return failure();
    }
    if (!matchPattern(producer, m_Op<arith::UIToFPOp>())) {
      return failure();
    }
  }

  // Producer of arith.uitofp op is arith.extui
  {
    producerOutput = producer->getOperand(0);
    producer = producerOutput.getDefiningOp();
    if (!producer) {
      return failure();
    }
    if (!matchPattern(producer, m_Op<arith::ExtUIOp>())) {
      return failure();
    }
  }

  // Ensure that the dequantization increases the
  // bitwidth from the input to the output
  auto elementTypeOut =
      cast<ShapedType>(genericOp.getOutputs()[0].getType()).getElementType();
  if (!elementTypeOut.isIntOrFloat()) {
    return failure();
  }
  unsigned bitWidthOut = elementTypeOut.getIntOrFloatBitWidth();
  auto elementTypeIn =
      cast<ShapedType>(genericOp.getInputs()[0].getType()).getElementType();
  if (!elementTypeIn.isIntOrFloat()) {
    return failure();
  }
  unsigned bitWidthIn = elementTypeIn.getIntOrFloatBitWidth();
  if (bitWidthIn >= bitWidthOut) {
    return failure();
  }

  return success();
}

static SmallVector<utils::IteratorType>
getParallelAndReductionIterators(unsigned nLoops, unsigned nReduction) {
  SmallVector<utils::IteratorType> res(nLoops - nReduction,
                                       utils::IteratorType::parallel);
  res.append(nReduction, utils::IteratorType::reduction);
  return res;
}

struct QuantizedMatmulRewriter {
  QuantizedMatmulRewriter(RewriterBase &rewriter, linalg::GenericOp dequant,
                          linalg::GenericOp matmul, int quantizedBitWidth);
  std::optional<SmallVector<OpOperand *>> getDequantMatmulInputs();
  std::pair<SmallVector<AffineMap>, SmallVector<utils::IteratorType>>
  getGroupReductionMapsAndIterators(OpOperand *inputOperand);
  Value getGroupReductionInit(Value input);
  Value generateGroupMaxGeneric();
  Value generateScalesGeneric(Value groupMax);
  Value generateGroupSumsGeneric();
  SmallVector<Value> generateQuantizationGenerics();
  linalg::GenericOp
  generateQuantizedMatmulGeneric(SmallVector<Value> quantizeResults);
  linalg::GenericOp
  generateReassociatedDequantizationGeneric(SmallVector<Value> quantizeResults,
                                            Value quantizedIntegerMatmul);
  LogicalResult precondition();

private:
  // rewriter
  RewriterBase &rewriter;
  // linalg::GenericOp that performs the dequantization
  linalg::GenericOp dequant;
  // linalg::GenericOp that performs the matmul
  linalg::GenericOp matmul;
  // Destination bit width for dynamic quantization of the unquantized input
  // (`ins[1]`)
  int quantizedBitWidth;
  // Type for accumulation of integer matmul
  IntegerType accType;
  // Type for multiplication in integer matmul (should probably be same as
  // `accType`)
  IntegerType mulType;
  // Type for dynamic quantization of unquantized input
  IntegerType quantType;
  // inputs to the `dequant` and `matmul` ops
  // ins[0] = quantized matrix input to `dequant`
  // ins[1] = unquantized input to `matmul`
  // ins[2] = scales input to `dequant
  // ins[3] = zero points input to `dequant`
  // ins[4] = dequantized input to `matmul`
  SmallVector<OpOperand *> ins;
  // Location for rewrite
  Location loc;
};

// Takes as input the dequantization `linalg.generic` op and the matmul
// `linalg.generic` op, and returns the scales, zero points, quantized
// input matrix, unquantizaed input matrix, and dequantized result
// matrix.
// TODO(#) Have stricter matching on inputs. There may be cases where
// the current matching fails
std::optional<SmallVector<OpOperand *>>
QuantizedMatmulRewriter::getDequantMatmulInputs() {
  assert(!failed(isContractionWithTwoReductions(matmul)) &&
         "expected `matmul` to be a contraction with two reduction dimensions");
  assert(!failed(isGroupedDequantizationOp(dequant)) &&
         "expected `dequant` to be a grouped dequantization");
  OpOperand *scales, *zps, *quantMat, *unquantMat, *dequantMat;
  auto maps = dequant.getIndexingMapsArray();
  for (auto [idx, map] : enumerate(ArrayRef<AffineMap>(maps).drop_back())) {
    if (map.isIdentity()) {
      quantMat = dequant.getDpsInputOperand(idx);
    } else if (map.isProjectedPermutation(true)) {
      for (Operation &bodyOp : dequant.getBlock()->getOperations()) {
        if (isa<arith::MulFOp>(bodyOp)) {
          if (bodyOp.getOperand(1) == dequant.getBlock()->getArgument(idx)) {
            scales = dequant.getDpsInputOperand(idx);
            break;
          }
        } else if (isa<arith::SubFOp>(bodyOp)) {
          if (bodyOp.getOperand(1) == dequant.getBlock()->getArgument(idx)) {
            zps = dequant.getDpsInputOperand(idx);
            break;
          }
        }
      }
    }
  }
  Value dequantOut = dequant.getResult(0);
  if (matmul.getDpsInputOperand(0)->get() == dequantOut) {
    unquantMat = matmul.getDpsInputOperand(1);
    dequantMat = matmul.getDpsInputOperand(0);
  } else {
    unquantMat = matmul.getDpsInputOperand(0);
    dequantMat = matmul.getDpsInputOperand(1);
  }
  if (scales && zps && quantMat && unquantMat) {
    return SmallVector<OpOperand *>(
        {quantMat, unquantMat, scales, zps, dequantMat});
  }
  return std::nullopt;
}

QuantizedMatmulRewriter::QuantizedMatmulRewriter(RewriterBase &rewriter,
                                                 linalg::GenericOp dequant,
                                                 linalg::GenericOp matmul,
                                                 int quantizedBitWidth)
    : rewriter(rewriter), dequant(dequant), matmul(matmul),
      quantizedBitWidth(quantizedBitWidth), loc(dequant.getLoc()) {
  accType = rewriter.getI32Type();
  mulType = rewriter.getI32Type();
  quantType = rewriter.getIntegerType(quantizedBitWidth);
  std::optional<SmallVector<OpOperand *>> inputs = getDequantMatmulInputs();
  if (inputs) {
    ins = *inputs;
  }
}

LogicalResult QuantizedMatmulRewriter::precondition() {
  if (ins.size() != 5) {
    return rewriter.notifyMatchFailure(
        matmul,
        "expected `ins` to have 5 inputs for quantized matmul reassociation");
  }
  OpOperand *unquantizedInputOperand = ins[1];
  Value unquantizedInput = ins[1]->get();
  RankedTensorType unquantizedInputType =
      cast<RankedTensorType>(unquantizedInput.getType());
  SmallVector<int64_t> unquantizedInputShape(unquantizedInputType.getShape());
  AffineMap indexingMap =
      matmul.getMatchingIndexingMap(unquantizedInputOperand);
  SmallVector<utils::IteratorType> matmulIteratorTypes =
      matmul.getIteratorTypesArray();
  if (unquantizedInputShape.size() < 2) {
    return rewriter.notifyMatchFailure(
        matmul, "input expected to have a rank of at least 2");
  }
  if (matmulIteratorTypes[indexingMap.getNumDims() - 1] ==
          utils::IteratorType::parallel ||
      matmulIteratorTypes[indexingMap.getNumDims() - 2] ==
          utils::IteratorType::parallel) {
    return rewriter.notifyMatchFailure(
        matmul, "inner 2 dimensions of matmul expected to be reduction");
  }
  auto affineExprs = indexingMap.getResults();
  auto innerDim0 = dyn_cast<AffineDimExpr>(affineExprs.back());
  auto innerDim1 = dyn_cast<AffineDimExpr>(affineExprs[affineExprs.size() - 2]);
  if (!innerDim0 || !innerDim1 ||
      innerDim0.getPosition() != indexingMap.getNumDims() - 1 ||
      innerDim1.getPosition() != indexingMap.getNumDims() - 2) {
    return rewriter.notifyMatchFailure(
        matmul, "inner shape of input expected to be reduced in matmul");
  }
  Value scales = ins[2]->get();
  Value zps = ins[3]->get();
  if (!isa<FloatType>(unquantizedInputType.getElementType()) ||
      !isa<FloatType>(getElementTypeOrSelf(scales)) ||
      !isa<FloatType>(getElementTypeOrSelf(zps))) {
    return rewriter.notifyMatchFailure(matmul, "expected float type");
  }
  OpOperand *matmulDequantizedOperand = ins[4];
  auto matmulDequantizedInputExprs =
      matmul.getMatchingIndexingMap(matmulDequantizedOperand).getResults();
  auto scalesType = dyn_cast<RankedTensorType>(scales.getType());
  auto zpsType = dyn_cast<RankedTensorType>(zps.getType());
  if (!scalesType || !zpsType) {
    return rewriter.notifyMatchFailure(
        dequant, "expected scales and zero points to have RankedTensorType");
  }
  if (scalesType.getShape().size() != matmulDequantizedInputExprs.size() - 1) {
    if (scalesType.getShape().size() != matmulDequantizedInputExprs.size() ||
        scalesType.getShape().back() != 1) {
      return rewriter.notifyMatchFailure(dequant, "unexpected rank for scales");
    }
  }
  if (zpsType.getShape().size() != matmulDequantizedInputExprs.size() - 1) {
    if (zpsType.getShape().size() != matmulDequantizedInputExprs.size() ||
        zpsType.getShape().back() != 1) {
      return rewriter.notifyMatchFailure(dequant,
                                         "unexpected rank for zero points");
    }
  }
  return success();
}

std::pair<SmallVector<AffineMap>, SmallVector<utils::IteratorType>>
QuantizedMatmulRewriter::getGroupReductionMapsAndIterators(
    OpOperand *inputOperand) {
  Value input = inputOperand->get();
  RankedTensorType inputType = cast<RankedTensorType>(input.getType());
  SmallVector<int64_t> inputShape(inputType.getShape());
  AffineMap indexingMap = matmul.getMatchingIndexingMap(inputOperand);
  SmallVector<utils::IteratorType> matmulIteratorTypes =
      matmul.getIteratorTypesArray();
  assert(inputShape.size() >= 2 &&
         "input expected to have a rank of at least 2");
  assert((matmulIteratorTypes[indexingMap.getNumDims() - 1] ==
              utils::IteratorType::reduction &&
          matmulIteratorTypes[indexingMap.getNumDims() - 2] ==
              utils::IteratorType::reduction) &&
         "inner 2 dimensions of matmul expected to be reduction");
  auto affineExprs = indexingMap.getResults();
  auto innerDim0 = dyn_cast<AffineDimExpr>(affineExprs.back());
  auto innerDim1 = dyn_cast<AffineDimExpr>(affineExprs[affineExprs.size() - 2]);
  assert(innerDim0 && innerDim1 &&
         innerDim0.getPosition() == indexingMap.getNumDims() - 1 &&
         innerDim1.getPosition() == indexingMap.getNumDims() - 2 &&
         "inner shape of input expected to be reduced in matmul");
  (void)innerDim0;
  (void)innerDim1;

  SmallVector<utils::IteratorType> iterators(inputShape.size(),
                                             utils::IteratorType::parallel);
  iterators.back() = utils::IteratorType::reduction;
  AffineMap inputMap = rewriter.getMultiDimIdentityMap(inputShape.size());
  AffineMap outputMap = rewriter.getMultiDimIdentityMap(inputShape.size())
                            .getMajorSubMap(inputShape.size() - 1);
  SmallVector<AffineMap> maps{inputMap, outputMap};
  return std::make_pair(maps, iterators);
}

// Helper to create an init Value for reductions along the group dimension.
Value QuantizedMatmulRewriter::getGroupReductionInit(Value input) {
  RankedTensorType inputType = cast<RankedTensorType>(input.getType());
  assert(isa<FloatType>(inputType.getElementType()) && "expected float type");
  Value zero = arith::ConstantOp::create(
      rewriter, loc, rewriter.getFloatAttr(inputType.getElementType(), 0.0));
  SmallVector<int64_t> inputShape(inputType.getShape());
  SmallVector<int64_t> outputShape(llvm::drop_end(inputShape));
  RankedTensorType outputType =
      RankedTensorType::get(outputShape, inputType.getElementType());
  Value emptyOut = tensor::EmptyOp::create(rewriter, loc, outputType.getShape(),
                                           outputType.getElementType());
  return linalg::FillOp::create(rewriter, loc, zero, emptyOut).result();
}

// Creates a generic that computes the absolute max along the group
// dimension, and returns the result.
Value QuantizedMatmulRewriter::generateGroupMaxGeneric() {
  OpOperand *inputOperand = ins[1];
  std::pair<SmallVector<AffineMap>, SmallVector<utils::IteratorType>>
      mapsAndIterators = getGroupReductionMapsAndIterators(inputOperand);
  auto maps = mapsAndIterators.first;
  auto iterators = mapsAndIterators.second;
  Value input = inputOperand->get();
  Value output = getGroupReductionInit(input);
  auto groupMaxOp = linalg::GenericOp::create(
      rewriter, loc, output.getType(), input, output, maps, iterators,
      [&](OpBuilder &b, Location nestedLoc, ValueRange args) {
        Value abs = math::AbsFOp::create(b, nestedLoc, args[0]);
        Value max = arith::MaximumFOp::create(b, nestedLoc, abs, args[1]);
        linalg::YieldOp::create(b, nestedLoc, max);
      });
  LLVM_DEBUG(DBGS() << "groupMaxOp:   " << groupMaxOp << "\n");
  return groupMaxOp.getResult(0);
}

// Creates a generic that computes the scales for each group, and
// returns the result.
Value QuantizedMatmulRewriter::generateScalesGeneric(Value groupMax) {
  auto groupMaxType = cast<RankedTensorType>(groupMax.getType());
  assert(isa<FloatType>(groupMaxType.getElementType()) &&
         "expected float type");
  Value cst = arith::ConstantOp::create(
      rewriter, loc,
      rewriter.getFloatAttr(groupMaxType.getElementType(),
                            (1 << (quantizedBitWidth - 1)) - 1));
  Value output = tensor::EmptyOp::create(rewriter, loc, groupMaxType.getShape(),
                                         groupMaxType.getElementType());
  SmallVector<AffineMap> maps(
      2, rewriter.getMultiDimIdentityMap(groupMaxType.getShape().size()));
  auto scalesOp = linalg::GenericOp::create(
      rewriter, loc, output.getType(), groupMax, output, maps,
      getParallelAndReductionIterators(groupMaxType.getRank(), 0),
      [&](OpBuilder &b, Location nestedLoc, ValueRange args) {
        Value scale = arith::DivFOp::create(b, nestedLoc, args[0], cst);
        linalg::YieldOp::create(b, nestedLoc, scale);
      });
  LLVM_DEBUG(DBGS() << "scalesOp:   " << scalesOp << "\n");
  return scalesOp.getResult(0);
}

// Creates a generic that computes the sums for each group, and
// returns the result.
Value QuantizedMatmulRewriter::generateGroupSumsGeneric() {
  OpOperand *inputOperand = ins[1];
  std::pair<SmallVector<AffineMap>, SmallVector<utils::IteratorType>>
      mapsAndIterators = getGroupReductionMapsAndIterators(inputOperand);
  auto maps = mapsAndIterators.first;
  auto iterators = mapsAndIterators.second;
  Value input = inputOperand->get();
  Value output = getGroupReductionInit(input);
  auto groupSumsOp = linalg::GenericOp::create(
      rewriter, loc, output.getType(), input, output, maps, iterators,
      [&](OpBuilder &b, Location nestedLoc, ValueRange args) {
        Value sum = arith::AddFOp::create(b, nestedLoc, args[0], args[1]);
        linalg::YieldOp::create(b, nestedLoc, sum);
      });
  LLVM_DEBUG(DBGS() << "groupSumsOp:   " << groupSumsOp << "\n");
  return groupSumsOp.getResult(0);
}

// Creates 4 linalg::GenericOps that collectively perform a symmetric
// quantization of the unquantized input, and returns the results.
SmallVector<Value> QuantizedMatmulRewriter::generateQuantizationGenerics() {
  assert(ins.size() == 5 && "expected `ins` to have 5 inputs");
  OpOperand *unquantizedInputOperand = ins[1];
  Value unquantizedInput = unquantizedInputOperand->get();

  auto unquantizedType = cast<RankedTensorType>(unquantizedInput.getType());
  SmallVector<int64_t> unquantizedShape(unquantizedType.getShape());

  IntegerType quantizedElementType = rewriter.getIntegerType(quantizedBitWidth);
  Value groupMax = generateGroupMaxGeneric();
  Value scales = generateScalesGeneric(groupMax);
  Value groupSums = generateGroupSumsGeneric();

  Value output = tensor::EmptyOp::create(rewriter, loc, unquantizedShape,
                                         quantizedElementType);
  AffineMap inputMap = rewriter.getMultiDimIdentityMap(unquantizedShape.size());
  AffineMap scalesMap = rewriter.getMultiDimIdentityMap(unquantizedShape.size())
                            .getMajorSubMap(unquantizedShape.size() - 1);
  AffineMap outputMap =
      rewriter.getMultiDimIdentityMap(unquantizedShape.size());
  SmallVector<AffineMap> maps{inputMap, scalesMap, outputMap};
  auto quantizeOp = linalg::GenericOp::create(
      rewriter, loc, output.getType(), ValueRange{unquantizedInput, scales},
      output, maps,
      getParallelAndReductionIterators(unquantizedShape.size(), 0),
      [&](OpBuilder &b, Location nestedLoc, ValueRange args) {
        Value scaled = arith::DivFOp::create(b, nestedLoc, args[0], args[1]);
        Value quant =
            arith::FPToSIOp::create(b, nestedLoc, quantizedElementType, scaled);
        linalg::YieldOp::create(b, nestedLoc, quant);
      });
  LLVM_DEBUG(DBGS() << "quantizeOp:   " << quantizeOp << "\n");
  Value newQuantizedInput = quantizeOp.getResult(0);
  SmallVector<Value> results{groupMax, scales, groupSums, newQuantizedInput};
  return results;
}

// Creates a generic that computes the main matmul computation in integer
// arithmetic, while values are still quantized, and returns the result.
linalg::GenericOp QuantizedMatmulRewriter::generateQuantizedMatmulGeneric(
    SmallVector<Value> quantizeResults) {
  Value newQuantizedInput = quantizeResults.back();
  assert(ins.size() == 5 && "expected `ins` to have 5 inputs");
  Value quantizedInput = ins[0]->get();
  OpOperand *matmulUnquantizedOperand = ins[1];
  OpOperand *matmulDequantizedOperand = ins[4];
  OpOperand *matmulOutputOperand = matmul.getDpsInitOperand(0);

  SmallVector<utils::IteratorType> iterators =
      getParallelAndReductionIterators(matmul.getNumLoops(), 1);
  SmallVector<AffineMap> maps;
  AffineMap matmulUnquantizedMap =
      matmul.getMatchingIndexingMap(matmulUnquantizedOperand);
  AffineMap matmulDequantizedMap =
      matmul.getMatchingIndexingMap(matmulDequantizedOperand);
  AffineMap matmulOutputMap =
      matmul.getMatchingIndexingMap(matmulOutputOperand);
  maps.push_back(matmulUnquantizedMap);
  maps.push_back(matmulDequantizedMap);
  SmallVector<AffineExpr> outputExprs(matmulOutputMap.getResults());
  outputExprs.push_back(rewriter.getAffineDimExpr(matmul.getNumLoops() - 2));
  maps.push_back(AffineMap::get(iterators.size(), 0, outputExprs,
                                outputExprs.front().getContext()));

  SmallVector<int64_t> newQuantizedInputShape(
      cast<RankedTensorType>(newQuantizedInput.getType()).getShape());
  assert(newQuantizedInputShape.size() >= 2 &&
         "expected new quantized input to have a rank of at least 2");
  SmallVector<int64_t> outputShape(
      cast<RankedTensorType>(matmulOutputOperand->get().getType()).getShape());
  outputShape.push_back(
      newQuantizedInputShape[newQuantizedInputShape.size() - 2]);
  Value zero = arith::ConstantOp::create(rewriter, loc,
                                         rewriter.getIntegerAttr(accType, 0.0));
  Value emptyOut = tensor::EmptyOp::create(rewriter, loc, outputShape, accType);
  Value output = linalg::FillOp::create(rewriter, loc, zero, emptyOut).result();
  auto integerMatmulOp = linalg::GenericOp::create(
      rewriter, loc, output.getType(),
      ValueRange{newQuantizedInput, quantizedInput}, output, maps, iterators,
      [&](OpBuilder &b, Location nestedLoc, ValueRange args) {
        Value mul;
        if (quantType == mulType) {
          Value ext1 = arith::ExtUIOp::create(b, nestedLoc, mulType, args[1]);
          mul = arith::MulIOp::create(b, nestedLoc, args[0], ext1);
        } else {
          Value ext0 = arith::ExtSIOp::create(b, nestedLoc, mulType, args[0]);
          Value ext1 = arith::ExtUIOp::create(b, nestedLoc, mulType, args[1]);
          mul = arith::MulIOp::create(b, nestedLoc, ext0, ext1);
        }
        Value sum;
        if (mulType == accType) {
          sum = arith::AddIOp::create(b, nestedLoc, mul, args[2]);
        } else {
          Value extMul = arith::ExtSIOp::create(b, nestedLoc, accType, mul);
          sum = arith::AddIOp::create(b, nestedLoc, extMul, args[2]);
        }
        linalg::YieldOp::create(b, nestedLoc, sum);
      });
  LLVM_DEBUG(DBGS() << "integerMatmulOp:   " << integerMatmulOp << "\n");
  return integerMatmulOp;
}

// Creates a generic that does the reassociated dequantization math, and does
// the final reduction along the number of groups dimension.
linalg::GenericOp
QuantizedMatmulRewriter::generateReassociatedDequantizationGeneric(
    SmallVector<Value> quantizeResults, Value quantizedIntegerMatmul) {
  assert(quantizeResults.size() == 4 &&
         "expected 4 ops from quantization step");
  Value scales = ins[2]->get();
  Value zps = ins[3]->get();
  Value newScales = quantizeResults[1];
  Value groupSums = quantizeResults[2];

  SmallVector<utils::IteratorType> iterators =
      getParallelAndReductionIterators(matmul.getNumLoops() - 1, 1);
  SmallVector<AffineMap> maps;
  AffineMap quantizedIntegerMatmulMap =
      rewriter.getMultiDimIdentityMap(iterators.size());
  maps.push_back(quantizedIntegerMatmulMap);

  {
    OpOperand *matmulUnquantizedOperand = ins[1];
    auto exprs =
        matmul.getMatchingIndexingMap(matmulUnquantizedOperand).getResults();
    exprs = exprs.slice(0, exprs.size() - 1);
    auto newScalesAndSumsMap =
        AffineMap::get(iterators.size(), 0, exprs, exprs.front().getContext());
    maps.push_back(newScalesAndSumsMap);
    maps.push_back(newScalesAndSumsMap);
  }

  {
    OpOperand *matmulDequantizedOperand = ins[4];
    auto exprsRef =
        matmul.getMatchingIndexingMap(matmulDequantizedOperand).getResults();
    SmallVector<AffineExpr> exprs(exprsRef.slice(0, exprsRef.size() - 1));
    RankedTensorType scalesType = cast<RankedTensorType>(scales.getType());
    RankedTensorType zpsType = cast<RankedTensorType>(zps.getType());
    if (exprs.size() < scalesType.getShape().size() &&
        scalesType.getShape().back() == 1 && zpsType.getShape().back() == 1) {
      exprs.push_back(rewriter.getAffineConstantExpr(0));
    }
    assert(exprs.size() == scalesType.getShape().size() &&
           "unexpected rank for scales");
    assert(exprs.size() == zpsType.getShape().size() &&
           "unexpected rank for zero points");
    auto scalesAndZpsMap =
        AffineMap::get(iterators.size(), 0, exprs, exprs.front().getContext());
    maps.push_back(scalesAndZpsMap);
    maps.push_back(scalesAndZpsMap);
  }

  OpOperand *matmulOutputOperand = matmul.getDpsInitOperand(0);
  SmallVector<AffineExpr> outputExprs(
      matmul.getMatchingIndexingMap(matmulOutputOperand).getResults());
  auto outputMap = AffineMap::get(iterators.size(), 0, outputExprs,
                                  outputExprs.front().getContext());
  maps.push_back(outputMap);

  Type floatType = getElementTypeOrSelf(scales);
  Value output = matmulOutputOperand->get();
  auto reassociatedDequantizationOp = linalg::GenericOp::create(
      rewriter, loc, output.getType(),
      ValueRange{quantizedIntegerMatmul, newScales, groupSums, scales, zps},
      output, maps, iterators,
      [&](OpBuilder &b, Location loc, ValueRange args) {
        Value dq;
        dq = arith::SIToFPOp::create(b, loc, floatType, args[0]);
        Value scaledRes0 = arith::MulFOp::create(b, loc, dq, args[1]);
        Value scaledRes1 = arith::MulFOp::create(b, loc, scaledRes0, args[3]);
        Value scaledZp0 = arith::MulFOp::create(b, loc, args[4], args[3]);
        Value scaledZp1 = arith::MulFOp::create(b, loc, scaledZp0, args[2]);
        Value groupRes = arith::SubFOp::create(b, loc, scaledRes1, scaledZp1);
        Value sum = arith::AddFOp::create(b, loc, groupRes, args[5]);
        linalg::YieldOp::create(b, loc, sum);
      });
  LLVM_DEBUG(DBGS() << "reassociatedDequantizationOp:   "
                    << reassociatedDequantizationOp << "\n");
  return reassociatedDequantizationOp;
}

// This function does the bulk of the rewrite for the dequantization + matmul.
//
// Starting with 2 `linalg.generic` ops (dequantization->matmul)
// %arg0 = quantized input
// %arg1 = scales
// %arg2 = zero points
// %arg3 = unquantized input
// ```mlir
// %0 = linalg.generic ins(%arg0, %arg1, %arg2 : tensor<8x4x2xi4>,
// tensor<8x4x1xf32>, tensor<8x4x1xf32>) outs(%1 : tensor<8x4x2xf32>) {
// ^bb0(%in: i4, %in_0: f32, %in_1: f32, %out: f32):
//   %9 = arith.extui %in : i4 to i32
//   %10 = arith.uitofp %9 : i32 to f32
//   %11 = arith.subf %10, %in_1 : f32
//   %12 = arith.mulf %11, %in_0 : f32
//   linalg.yield %12 : f32
// } -> tensor<8x4x2xf32>
// %2 = linalg.generic ins(%arg3, %0 : tensor<4x2xf32>, tensor<8x4x2xf32>)
// outs(%3 : tensor<8xf32>) { ^bb0(%in: f32, %in_0: f32, %out: f32):
//   %9 = arith.mulf %in, %in_0 : f32
//   %10 = arith.addf %9, %out : f32
//   linalg.yield %10 : f32
// } -> tensor<8xf32>
// ```
//
// This function rewrites the above ops as the following new sequence of 6 ops
// that does the following:
//
// a) Dynamically quantize the unquantized input:
//      1. Compute the absolute max of the unquantized input (%arg3) within each
//      group.
//      2. Compute scales for %arg3 by dividing the absolute max by (1 <<
//      newBitWidth) - 1),
//         where newBitWidth is the bitwidth of the new quantized type,
//         currently set to `i16`.
//      3. Compute the sum along groups of the unquantized input. This is not
//      necessary for
//         the quantization step, but it is needed to efficiently perform a
//         reassociated quantized matmul in steps 5-6.
//      4. Quantize the unquantized input (%arg3) by dividing elements in each
//      group by
//         the corresponding scale. This quantization is symmetric with no zero
//         point.
// b) Perform the reassociated quantized matmul, keeping the bulk of the
// computation in
//    integer arithmetic:
//      5. The first op performs a matmul-like operation that reduces the
//      innermost group
//         dimension. Note the indexing maps in the following example:
//         ```mlir
//         %22 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) ->
//         (d1, d2)>,
//                                                affine_map<(d0, d1, d2) ->
//                                                (d0, d1, d2)>, affine_map<(d0,
//                                                d1, d2) -> (d0, d1)>],
//                               iterator_types = ["parallel", "parallel",
//                               "reduction"]} ins(%17, %0 : tensor<4x2xi16>,
//                               tensor<8x4x2xi4>) outs(%19 : tensor<8x4xi32>) {
//             ^bb0(%in: i16, %in_4: i4, %out: i32):
//               %24 = arith.extsi %in : i16 to i32
//               %25 = arith.extui %in_4 : i4 to i32
//               %26 = arith.muli %24, %25 : i32
//               %27 = arith.addi %26, %out : i32
//               linalg.yield %27 : i32
//             } -> tensor<8x4xi32>
//         ```
//         This op also extends the inputs to the accumulation type, i32 in this
//         case, to target specific x86 instructions. We perform the matrix
//         multiplication before the dequantization arithmetic, which has been
//         reassociated into op 6.
//      6. The final op performs the remaining reduction across groups and does
//      the
//         dequantization arithmetic:
//         ```mlir
//         %23 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0,
//         d1)>,
//                                                affine_map<(d0, d1) -> (d1)>,
//                                                affine_map<(d0, d1) -> (d1)>,
//                                                affine_map<(d0, d1) -> (d0,
//                                                d1)>, affine_map<(d0, d1) ->
//                                                (d0, d1)>, affine_map<(d0, d1)
//                                                -> (d0)>],
//                               iterator_types = ["parallel", "reduction"]}
//                               ins(tensor<8x4xi32>, tensor<4xf32>,
//                               tensor<4xf32>, tensor<8x4xf32>,
//                               tensor<8x4xf32>) outs(tensor<8xf32>) {
//             ^bb0(%in: i32, %in_4: f32, %in_5: f32, %in_6: f32, %in_7: f32,
//             %out: f32):
//               %24 = arith.sitofp %in : i32 to f32
//               %25 = arith.mulf %24, %in_4 : f32
//               %26 = arith.mulf %25, %in_6 : f32
//               %27 = arith.mulf %in_7, %in_6 : f32
//               %28 = arith.mulf %27, %in_5 : f32
//               %29 = arith.subf %26, %28 : f32
//               %30 = arith.addf %29, %out : f32
//               linalg.yield %30 : f32
//             } -> tensor<8xf32>
//         ```
//
// ** Note that this rewrite introduces precision loss in the matmul, and is a
//    tradeoff between precision and performance. This rewrite should most
//    likely be opt-in only. **
static LogicalResult reassociateDequantMatmul(RewriterBase &rewriter,
                                              linalg::GenericOp dequant,
                                              linalg::GenericOp matmul,
                                              int quantizeBitWidth) {
  QuantizedMatmulRewriter qmr(rewriter, dequant, matmul, quantizeBitWidth);
  if (failed(qmr.precondition())) {
    return success();
  }
  SmallVector<Value> quantizeResults = qmr.generateQuantizationGenerics();
  linalg::GenericOp quantizedIntegerMatmul =
      qmr.generateQuantizedMatmulGeneric(quantizeResults);
  linalg::GenericOp reassociatedDequantization =
      qmr.generateReassociatedDequantizationGeneric(
          quantizeResults, quantizedIntegerMatmul.getResult(0));

  rewriter.replaceOp(matmul, reassociatedDequantization.getResult(0));

  return success();
}

struct FuseDequantizationMatmulPass
    : impl::FuseDequantizationMatmulPassBase<FuseDequantizationMatmulPass> {
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<linalg::LinalgDialect, IREE::Flow::FlowDialect,
                    math::MathDialect>();
  }
  void runOnOperation() override;
};

} // namespace

void FuseDequantizationMatmulPass::runOnOperation() {
  MLIRContext *context = &getContext();
  mlir::FunctionOpInterface funcOp = getOperation();

  int quantizeBitWidth = 16;
  SmallVector<std::pair<linalg::GenericOp, linalg::GenericOp>> candidates;
  for (auto genericOp : funcOp.getFunctionBody().getOps<linalg::GenericOp>()) {
    if (failed(isContractionWithTwoReductions(genericOp))) {
      continue;
    }

    OpOperand *lhs = genericOp.getDpsInputOperand(0);
    OpOperand *rhs = genericOp.getDpsInputOperand(1);
    auto lhsOp = lhs->get().getDefiningOp<linalg::GenericOp>();
    auto rhsOp = rhs->get().getDefiningOp<linalg::GenericOp>();
    if (!cast<ShapedType>(genericOp.getInputs()[0].getType())
             .hasStaticShape() ||
        !cast<ShapedType>(genericOp.getInputs()[1].getType())
             .hasStaticShape() ||
        !cast<ShapedType>(genericOp.getResults()[0].getType())
             .hasStaticShape()) {
      // Codegen can't handle the dynamic case yet.
      continue;
    }
    if (lhsOp) {
      if (!failed(isGroupedDequantizationOp(lhsOp))) {
        candidates.push_back(std::make_pair(lhsOp, genericOp));
        continue;
      }
    }
    if (rhsOp) {
      if (!failed(isGroupedDequantizationOp(rhsOp))) {
        candidates.push_back(std::make_pair(rhsOp, genericOp));
      }
    }
  }
  IRRewriter rewriter(context);
  for (auto candidate : candidates) {
    rewriter.setInsertionPointAfter(candidate.second);
    if (failed(reassociateDequantMatmul(rewriter, candidate.first,
                                        candidate.second, quantizeBitWidth))) {
      return signalPassFailure();
    }
  }
}
} // namespace mlir::iree_compiler::GlobalOptimization
