// 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/Codegen/Common/EmulateNarrowType.h"
#include "iree/compiler/Codegen/Common/Transforms.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
#include "iree/compiler/Dialect/HAL/IR/HALDialect.h"
#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Transforms/NarrowTypeEmulationConverter.h"
#include "mlir/Dialect/Arith/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

namespace mlir::iree_compiler {

#define GEN_PASS_DEF_EMULATENARROWTYPEPASS
#include "iree/compiler/Codegen/Common/Passes.h.inc"

namespace {

//===----------------------------------------------------------------------===//
// Conversion patterns
//===----------------------------------------------------------------------===//

struct ConvertHalInterfaceBindingSubspan final
    : OpConversionPattern<IREE::HAL::InterfaceBindingSubspanOp> {
  using Base::Base;

  LogicalResult
  matchAndRewrite(IREE::HAL::InterfaceBindingSubspanOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto currentType = dyn_cast<MemRefType>(op.getType());
    if (!currentType) {
      return rewriter.notifyMatchFailure(op->getLoc(),
                                         "unhandled non-memref types");
    }
    auto newResultType =
        dyn_cast<MemRefType>(getTypeConverter()->convertType(currentType));
    if (!newResultType) {
      return rewriter.notifyMatchFailure(
          op->getLoc(),
          llvm::formatv("failed to legalize memref type: {}", op.getType()));
    }
    Location loc = op.getLoc();
    OpFoldResult zero = rewriter.getIndexAttr(0);
    SmallVector<OpFoldResult> indices(currentType.getRank(), zero);

    // Get linearized type.
    int srcBits = currentType.getElementType().getIntOrFloatBitWidth();
    int dstBits = newResultType.getElementType().getIntOrFloatBitWidth();
    OpFoldResult elementOffset;
    Value byteOffset = adaptor.getByteOffset();
    if (byteOffset && !matchPattern(byteOffset, m_Zero())) {
      elementOffset = convertByteOffsetToElementOffset(
          rewriter, loc, byteOffset, currentType.getElementType());
    } else {
      elementOffset = rewriter.getIndexAttr(0);
    }
    SmallVector<OpFoldResult> sizes = getMixedValues(
        currentType.getShape(), adaptor.getDynamicDims(), rewriter);
    memref::LinearizedMemRefInfo linearizedMemRefInfo =
        memref::getLinearizedMemRefOffsetAndSize(rewriter, loc, srcBits,
                                                 dstBits, elementOffset, sizes);

    SmallVector<Value> dynamicLinearizedSize;
    if (newResultType.getRank() > 0 && !newResultType.hasStaticShape()) {
      dynamicLinearizedSize.push_back(getValueOrCreateConstantIndexOp(
          rewriter, loc, linearizedMemRefInfo.linearizedSize));
    }

    rewriter.replaceOpWithNewOp<IREE::HAL::InterfaceBindingSubspanOp>(
        op, newResultType, adaptor.getLayout(), adaptor.getBinding(),
        byteOffset, dynamicLinearizedSize, adaptor.getAlignmentAttr(),
        adaptor.getDescriptorFlagsAttr());
    return success();
  }
};

static void populateIreeNarrowTypeEmulationPatterns(
    arith::NarrowTypeEmulationConverter &converter,
    RewritePatternSet &patterns) {
  patterns.add<ConvertHalInterfaceBindingSubspan>(converter,
                                                  patterns.getContext());
}

//===----------------------------------------------------------------------===//
// Pass Definition
//===----------------------------------------------------------------------===//

struct EmulateNarrowTypePass final
    : impl::EmulateNarrowTypePassBase<EmulateNarrowTypePass> {
  using Base::Base;

  void getDependentDialects(DialectRegistry &registry) const override {
    registry
        .insert<arith::ArithDialect, func::FuncDialect, memref::MemRefDialect,
                vector::VectorDialect, affine::AffineDialect,
                IREE::Codegen::IREECodegenDialect, IREE::HAL::HALDialect>();
  }

  void runOnOperation() override {
    if (failed(emulateNarrowType(getOperation(), disableAtomicRMW))) {
      return signalPassFailure();
    }
  }
};
} // namespace

LogicalResult emulateNarrowType(
    Operation *root, bool disableAtomicRMW,
    std::optional<NarrowTypeConversionPopulationFn> populateCallback) {
  // The number of bits used in a load/store op.
  constexpr unsigned kLoadStoreEmulateBitwidth = 8;
  static_assert(
      llvm::isPowerOf2_32(kLoadStoreEmulateBitwidth) &&
      "only power of 2 is supported for narrow type load/store emulation");

  MLIRContext *ctx = root->getContext();

  // Resolve memref.dim ops before emulation. This is needed because the
  // emulation linearizes memrefs, changing their rank and shape semantics.
  // Any memref.dim on a narrow-type memref must be traced back to its source
  // dynamic dimensions before we can safely emulate.
  {
    RewritePatternSet dimPatterns(ctx);
    memref::populateResolveRankedShapedTypeResultDimsPatterns(dimPatterns);
    GreedyRewriteConfig config;
    config.setRegionSimplificationLevel(GreedySimplifyRegionLevel::Disabled);
    if (failed(applyPatternsGreedily(root, std::move(dimPatterns), config))) {
      return root->emitOpError("failed to resolve shaped type result dims");
    }
  }

  arith::NarrowTypeEmulationConverter typeConverter(kLoadStoreEmulateBitwidth);
  memref::populateMemRefNarrowTypeEmulationConversions(typeConverter);

  ConversionTarget target(*ctx);
  target.addDynamicallyLegalOp<func::FuncOp>([&typeConverter](Operation *op) {
    return typeConverter.isLegal(cast<func::FuncOp>(op).getFunctionType());
  });
  auto opLegalCallback = [&typeConverter](Operation *op) {
    return typeConverter.isLegal(op);
  };
  target.addDynamicallyLegalOp<func::CallOp, func::ReturnOp>(opLegalCallback);
  target.addDynamicallyLegalDialect<
      arith::ArithDialect, vector::VectorDialect, memref::MemRefDialect,
      affine::AffineDialect, IREE::HAL::HALDialect>(opLegalCallback);

  RewritePatternSet patterns(ctx);

  // Try to flatten memrefs as a prerequisite for narrow type emulation,
  // so we can have simplified checks in the emulation patterns.
  memref::populateFlattenMemrefsPatterns(patterns);

  arith::populateArithNarrowTypeEmulationPatterns(typeConverter, patterns);
  memref::populateMemRefNarrowTypeEmulationPatterns(typeConverter, patterns,
                                                    disableAtomicRMW);
  populateIREEResolveExtractStridedMetadataPatterns(patterns);
  vector::populateVectorNarrowTypeEmulationPatterns(typeConverter, patterns,
                                                    disableAtomicRMW,
                                                    /*assumeAligned=*/true);
  populateIreeNarrowTypeEmulationPatterns(typeConverter, patterns);
  if (populateCallback) {
    populateCallback.value()(typeConverter, patterns, target);
  }

  if (failed(applyPartialConversion(root, target, std::move(patterns)))) {
    return root->emitOpError("failed to emulate bit width");
  }

  GreedyRewriteConfig config;
  config.setRegionSimplificationLevel(GreedySimplifyRegionLevel::Normal);

  RewritePatternSet sinkBroadcast(ctx);
  vector::populateSinkVectorOpsPatterns(sinkBroadcast);
  if (failed(applyPatternsGreedily(root, std::move(sinkBroadcast), config))) {
    return root->emitOpError("failed in sinking of broadcasts");
  }

  // Also do the `bitcast -> extui/extsi` rewrite.
  RewritePatternSet foldExtPatterns(ctx);
  vector::populateVectorNarrowTypeRewritePatterns(foldExtPatterns);
  if (failed(applyPatternsGreedily(root, std::move(foldExtPatterns), config))) {
    return failure();
  }
  return success();
}

} // namespace mlir::iree_compiler
