// Copyright 2024 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/GPU/Passes.h"
#include "iree/compiler/Codegen/Common/Transforms.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
#include "iree/compiler/Codegen/Dialect/GPU/Transforms/Transforms.h"
#include "iree/compiler/Codegen/Transforms/Transforms.h"
#include "iree/compiler/Codegen/Utils/GPUUtils.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "llvm/Support/DebugLog.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Interfaces/DestinationStyleOpInterface.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

#define DEBUG_TYPE "iree-codegen-gpu-fuse-and-hoist-parallel-loops"

namespace mlir::iree_compiler {

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

using namespace IREE::GPU;

namespace {
struct GPUFuseAndHoistParallelLoopsPass final
    : impl::GPUFuseAndHoistParallelLoopsPassBase<
          GPUFuseAndHoistParallelLoopsPass> {
  void runOnOperation() override;
};
} // namespace

static std::optional<int64_t> getStaticForallTripCount(scf::ForallOp forall) {
  // TODO: Handle non-normalized loops.
  if (!forall.isNormalized()) {
    return std::nullopt;
  }
  int64_t tripCount = 1;
  for (OpFoldResult ub : forall.getMixedUpperBound()) {
    std::optional<int64_t> maybeConstantUb = getConstantIntValue(ub);
    if (!maybeConstantUb) {
      return std::nullopt;
    }
    tripCount *= *maybeConstantUb;
  }
  return tripCount;
}

static bool forallTripCountMatchesWorkgroupSize(scf::ForallOp forallOp,
                                                int64_t flatWorkgroupSize) {
  std::optional<int64_t> maybeTripCount = getStaticForallTripCount(forallOp);
  if (!maybeTripCount) {
    return false;
  }

  // For lane mapped foralls we need to verify that it is contained within
  // a parent warp mapped op that combines to match the workggroup size.
  if (forallOpHasMappingType<IREE::GPU::LaneIdAttr>(forallOp)) {
    auto parentForall = forallOp->getParentOfType<scf::ForallOp>();
    if (!parentForall ||
        !forallOpHasMappingType<gpu::GPUWarpMappingAttr>(parentForall)) {
      return false;
    }

    std::optional<int64_t> maybeParentTripCount =
        getStaticForallTripCount(parentForall);
    if (!maybeParentTripCount) {
      return false;
    }

    return *maybeParentTripCount * *maybeTripCount == flatWorkgroupSize;
  }

  // All other loops must be mapped to threads to compare.
  if (!forallOpHasMappingType<gpu::GPUThreadMappingAttr>(forallOp)) {
    return false;
  }

  return *maybeTripCount == flatWorkgroupSize;
}
struct FuseForalls final : OpRewritePattern<scf::ForallOp> {
  using Base::Base;
  FuseForalls(MLIRContext *ctx, int64_t flatWorkgroupSize, PatternBenefit b = 1)
      : OpRewritePattern<scf::ForallOp>(ctx, b),
        flatWorkgroupSize(flatWorkgroupSize) {}
  LogicalResult matchAndRewrite(scf::ForallOp producerForall,
                                PatternRewriter &rewriter) const override {
    if (!producerForall->hasOneUse()) {
      return rewriter.notifyMatchFailure(producerForall,
                                         "multi-use producer forall");
    }

    SmallVector<Operation *> consumerChain;
    Operation *currProducer = *producerForall->user_begin();
    while (currProducer && currProducer->hasOneUse()) {
      consumerChain.push_back(currProducer);
      if (!isa<tensor::ExpandShapeOp, tensor::CollapseShapeOp>(currProducer)) {
        break;
      }
      currProducer = *currProducer->user_begin();
    }

    auto consumerForall = currProducer->getParentOfType<scf::ForallOp>();
    if (!consumerForall || !forallTripCountMatchesWorkgroupSize(
                               consumerForall, flatWorkgroupSize)) {
      return rewriter.notifyMatchFailure(
          producerForall,
          "no consumer forall with trip count matching workgroup size");
    }

    // TODO: Allow extracting multiple uses within the same consumer loop. Still
    // single producer single consumer loop, but multiple uses within the
    // consumer.
    if (!producerForall->hasOneUse()) {
      return failure();
    }

    return fuseForallIntoConsumer(rewriter, producerForall, consumerForall,
                                  consumerChain);
  }

private:
  int64_t flatWorkgroupSize;
};

struct FuseNestedLaneAndWarpForalls final : OpRewritePattern<scf::ForallOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(scf::ForallOp warpForallOp,
                                PatternRewriter &rewriter) const override {
    SmallVector<scf::ForallOp> innerForallOps(
        warpForallOp.getBody()->getOps<scf::ForallOp>());
    if (innerForallOps.size() != 1) {
      return failure();
    }
    scf::ForallOp laneForallOp = innerForallOps[0];
    FailureOr<scf::ForallOp> threadForallOp =
        fuseNestedLaneAndWarpForalls(rewriter, warpForallOp, laneForallOp);
    if (failed(threadForallOp)) {
      return failure();
    }
    rewriter.replaceOp(warpForallOp, *threadForallOp);
    return success();
  }
};

struct FuseTilableDestinationProducers final : OpRewritePattern<scf::ForallOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(scf::ForallOp forallOp,
                                PatternRewriter &rewriter) const override {
    TilingInterface tileableProducer;
    tensor::ExtractSliceOp sliceOp;
    for (auto iterArg : forallOp.getRegionIterArgs()) {
      for (auto user : iterArg.getUsers()) {
        sliceOp = dyn_cast<tensor::ExtractSliceOp>(user);
        if (sliceOp) {
          break;
        }
      }
      if (!sliceOp) {
        continue;
      }
      tileableProducer = forallOp.getTiedLoopInit(iterArg)
                             ->get()
                             .getDefiningOp<TilingInterface>();
      // Pad fusion is handled separately as we dont want zero slice guards that
      // happen by default.
      if (tileableProducer && isa<tensor::PadOp>(tileableProducer)) {
        tileableProducer = nullptr;
      }
      if (tileableProducer) {
        break;
      }
    }
    if (!tileableProducer) {
      return failure();
    }

    SmallVector<LoopLikeOpInterface> loops = {forallOp};
    rewriter.startOpModification(forallOp);
    std::optional<scf::SCFFuseProducerOfSliceResult> fusionResult =
        mlir::scf::tileAndFuseProducerOfSlice(rewriter, sliceOp, loops);
    if (!fusionResult) {
      return failure();
    }
    rewriter.finalizeOpModification(forallOp);
    return success();
  }
};

struct FuseUnitLoopDestination final : OpRewritePattern<scf::ForallOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(scf::ForallOp forallOp,
                                PatternRewriter &rewriter) const override {
    std::optional<int64_t> maybeTripCount = getStaticForallTripCount(forallOp);
    if (!maybeTripCount || *maybeTripCount != 1) {
      return rewriter.notifyMatchFailure(forallOp,
                                         "not a unit trip count loop");
    }
    DestinationStyleOpInterface dpsProducer;
    BlockArgument bodyArg;
    Value dpsResult;
    for (auto iterArg : forallOp.getRegionIterArgs()) {
      dpsResult = forallOp.getTiedLoopInit(iterArg)->get();
      bodyArg = iterArg;
      dpsProducer = dpsResult.getDefiningOp<DestinationStyleOpInterface>();
      if (dpsProducer) {
        break;
      }
    }
    if (!dpsProducer || !dpsProducer->hasOneUse()) {
      return rewriter.notifyMatchFailure(forallOp,
                                         "no single use DPS producer");
    }

    Operation *parallelInsert = nullptr;
    for (auto user : bodyArg.getUsers()) {
      if (isa<tensor::ParallelInsertSliceOp>(user)) {
        // This should be illegal but check anyway.
        if (parallelInsert) {
          return rewriter.notifyMatchFailure(forallOp, "multiple insert users");
        }
        parallelInsert = user;
      }
    }
    if (!parallelInsert) {
      return rewriter.notifyMatchFailure(
          forallOp, "destination not used by a parallel insert");
    }

    rewriter.startOpModification(forallOp);
    // Move the producer into the body of the forall loop.
    rewriter.moveOpBefore(dpsProducer, forallOp.getBody(),
                          forallOp.getBody()->begin());

    // Replace all uses of the region iter arg with the moved dps op.
    rewriter.replaceAllUsesExcept(bodyArg, dpsResult, parallelInsert);

    // Set the init operand of the forall op to the init operand of the
    // producer.
    int64_t dpsInitIndex = cast<OpResult>(dpsResult).getResultNumber();
    forallOp->setOperand(forallOp.getTiedOpOperand(bodyArg)->getOperandNumber(),
                         dpsProducer.getDpsInitOperand(dpsInitIndex)->get());

    // Finally replace the init operand of the moved producer with the region
    // iter arg.
    dpsProducer.setDpsInitOperand(dpsInitIndex, bodyArg);
    rewriter.finalizeOpModification(forallOp);
    return success();
  }
};

struct FuseTilableSliceProducers final
    : OpRewritePattern<tensor::ExtractSliceOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
                                PatternRewriter &rewriter) const override {
    if (sliceOp->use_empty()) {
      return failure();
    }
    auto tilableProducer = sliceOp.getSource().getDefiningOp<TilingInterface>();
    // Pad fusion is handled separately as we dont want zero slice guards that
    // happen by default.
    if (!tilableProducer || isa<tensor::PadOp>(tilableProducer)) {
      return failure();
    }

    auto parentForall = sliceOp->getParentOfType<scf::ForallOp>();
    if (!parentForall) {
      return failure();
    }

    auto producerParent = tilableProducer->getParentOfType<scf::ForallOp>();
    if (producerParent && producerParent == parentForall) {
      return failure();
    }

    SmallVector<LoopLikeOpInterface> loops = {parentForall};
    std::optional<scf::SCFFuseProducerOfSliceResult> fusionResult =
        mlir::scf::tileAndFuseProducerOfSlice(rewriter, sliceOp, loops);
    if (!fusionResult) {
      return failure();
    }
    return success();
  }
};

struct FuseCollapseShapeConsumers final
    : OpRewritePattern<tensor::CollapseShapeOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(tensor::CollapseShapeOp collapseOp,
                                PatternRewriter &rewriter) const override {
    auto forallOp = collapseOp.getSrc().getDefiningOp<scf::ForallOp>();
    if (!forallOp) {
      return rewriter.notifyMatchFailure(collapseOp, "No forall op producer");
    }

    if (failed(fuseCollapseShapeIntoProducerForall(rewriter, forallOp,
                                                   collapseOp))) {
      return failure();
    }
    return success();
  }
};

struct FuseExtractSliceConsumers final
    : OpRewritePattern<tensor::ExtractSliceOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(tensor::ExtractSliceOp extractSliceOp,
                                PatternRewriter &rewriter) const override {
    // Find the scf::ForallOp producer, and get the corresponding
    // tensor::ParallelInsertSliceOp.
    auto forallOp = extractSliceOp.getSource().getDefiningOp<scf::ForallOp>();
    if (!forallOp) {
      return rewriter.notifyMatchFailure(extractSliceOp,
                                         "No forall op producer");
    }

    if (failed(fuseExtractSliceIntoProducerForall(rewriter, forallOp,
                                                  extractSliceOp))) {
      return failure();
    }
    return success();
  }
};

void GPUFuseAndHoistParallelLoopsPass::runOnOperation() {
  MLIRContext *context = &getContext();

  FunctionOpInterface funcOp = getOperation();

  // Try to get the flat workgroup size if possible.
  std::optional<int64_t> maybeFlatWorkgroupSize = std::nullopt;
  if (std::optional<SmallVector<int64_t>> workgroupSize =
          getWorkgroupSize(funcOp)) {
    maybeFlatWorkgroupSize = llvm::product_of(*workgroupSize);
  }

  // First run the hoisting and fusion patterns.
  {
    RewritePatternSet patterns(context);
    // These two patterns are run to a fixed point, allowing fusion within
    // potentially nested loops, hoisting from said loops, and continued fusion.
    if (maybeFlatWorkgroupSize) {
      // Forall fusion requires knowing the workgroup size to verify the fusion
      // is valid. Without validation we risk putting barriers inside
      // conditioned regions (e.g. scf.if/for).
      patterns.add<FuseForalls>(context, *maybeFlatWorkgroupSize,
                                /*benefit=*/1);
    }
    populateFuseTilableForallConsumersPattern(patterns);
    tensor::populateMergeConsecutiveInsertExtractSlicePatterns(patterns);
    patterns.add<FuseNestedLaneAndWarpForalls>(context);
    populateForallLoopHoistingPattern(patterns);
    if (failed(applyPatternsGreedily(funcOp, std::move(patterns)))) {
      funcOp->emitOpError("failed to apply fusion + hoisting patterns (set 1)");
      return signalPassFailure();
    }
  }

  LDBG() << "After fusing and hoisting loops\n" << funcOp;

  // After hoisting parallel loops, try to fuse in any newly revealed consumers
  // and destinations.
  // TODO: Move the consumer fusion pattern to an explicit worklist rather than
  // using the GreedyPatternRewriter.
  {
    RewritePatternSet patterns(context);
    patterns.add<FuseTilableDestinationProducers>(context);
    patterns.add<FuseUnitLoopDestination>(context);
    populateFuseTilableForallConsumersPattern(patterns);
    patterns.add<FuseCollapseShapeConsumers>(context);
    patterns.add<FuseExtractSliceConsumers>(context);
    // Only swap expand_shape with extract_slice when they are in different
    // blocks. This acts as a fusion pattern for the expand_shape, which should
    // only apply when they are in different blocks (i.e., one inside a loop,
    // and one outside). Swapping in other cases can interfere with later
    // optimizations that fold the slice into consumer load operations.
    auto swapControlFn = [](OpOperand *operand) {
      Operation *producer = operand->get().getDefiningOp();
      Operation *consumer = operand->getOwner();
      return producer->getBlock() != consumer->getBlock();
    };
    populateSwapExtractWithExpandPattern(patterns, swapControlFn);
    tensor::populateFoldTensorEmptyPatterns(patterns);
    scf::ForallOp::getCanonicalizationPatterns(patterns, context);
    if (failed(applyPatternsGreedily(funcOp, std::move(patterns)))) {
      funcOp->emitOpError("failed to apply fusion + hoisting patterns (set 2)");
      return signalPassFailure();
    }
  }

  LDBG() << "After fusing new consumers\n" << funcOp;

  // Finally try to do any new producer fusions.
  {
    RewritePatternSet patterns(context);
    patterns.add<FuseTilableDestinationProducers>(context);
    patterns.add<FuseTilableSliceProducers>(context);
    tensor::populateFoldTensorEmptyPatterns(patterns);
    scf::ForallOp::getCanonicalizationPatterns(patterns, context);
    auto zeroSliceGuard = [](tensor::ExtractSliceOp) -> std::optional<bool> {
      // Do not use zero slice guard.
      return false;
    };
    patterns.add<linalg::ExtractSliceOfPadTensorSwapPattern>(context,
                                                             zeroSliceGuard);
    if (failed(applyPatternsGreedily(funcOp, std::move(patterns)))) {
      funcOp->emitOpError("failed to apply fusion + hoisting patterns (set 3)");
      return signalPassFailure();
    }
  }

  LDBG() << "After fusing new producers\n" << funcOp;
}

} // namespace mlir::iree_compiler
