// 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/Dialect/Flow/IR/FlowDialect.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
#include "iree/compiler/Dialect/Flow/Transforms/ConvertRegionToWorkgroups.h"
#include "iree/compiler/Dialect/Flow/Transforms/FormDispatchRegions.h"
#include "iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
#include "iree/compiler/Dialect/TensorExt/IR/TensorExtOps.h"
#include "iree/compiler/DispatchCreation/Passes.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Support/LLVM.h"

#define DEBUG_TYPE                                                             \
  "iree-dispatch-creation-materialize-default-workgroup-count-region"

namespace mlir::iree_compiler::DispatchCreation {

#define GEN_PASS_DEF_MATERIALIZEDEFAULTWORKGROUPCOUNTREGIONPASS
#include "iree/compiler/DispatchCreation/Passes.h.inc"

/// Creates the workgroup count region where the materialized computation
/// is derived as a program slice of the body of the dispatch. This method
/// - Computes the `workload` to use for the `workgroupsOp`, which are
///   derived from the values captured by the `workgroupsOp`.
/// - Populates the workgroup count region for this with the placeholder
///   op `flow.dispatch.workgroups_count_from_body_slice`. This op is
///   resolved in the backends into the actual workgroup count computation.
/// - To correlate back to the captured workload,
/// `iree_tensor_ext.dispatch.workload.ordinal`
///   to map the captured operand to the position in the workload list.
static LogicalResult createDefaultWorkgroupCountRegion(
    RewriterBase &rewriter, IREE::Flow::DispatchWorkgroupsOp workgroupsOp) {
  Region &workgroupCountBody = workgroupsOp.getWorkgroupCount();
  if (!workgroupCountBody.empty()) {
    // Preserve pre-existing workgroup count region.
    return success();
  }

  // Compute the `workload`. For now all `IndexType` are treated as workload.
  SmallVector<Value> workload;
  SmallVector<Type> workloadTypes;
  SmallVector<Location> workloadLocs;
  for (auto argument : workgroupsOp.getArguments()) {
    Type argumentType = argument.getType();
    if (!isa<IndexType>(argumentType)) {
      continue;
    }
    workload.push_back(argument);
    workloadTypes.push_back(argumentType);
    workloadLocs.push_back(argument.getLoc());
  }

  // Populate the count region.
  Block *block =
      rewriter.createBlock(&workgroupCountBody, workgroupCountBody.end(),
                           workloadTypes, workloadLocs);
  Location loc = workgroupsOp.getLoc();
  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPointToStart(block);
  Operation *defaultCountOp =
      IREE::TensorExt::DispatchWorkgroupCountFromSliceOp::create(
          rewriter, loc, block->getArguments());

  // Check for presence of `scf.forall` operations created by partial reduction
  // tiling.
  auto forallOps = workgroupsOp.getClosureBodyRegion().getOps<scf::ForallOp>();
  if (!forallOps.empty()) {
    if (!llvm::hasSingleElement(forallOps)) {
      return workgroupsOp->emitOpError(
          "unhandled multiple scf.forall ops in a dispatch");
    }
    auto forallOp = *forallOps.begin();
    std::optional<ArrayAttr> mapping = forallOp.getMapping();
    if (!mapping) {
      return workgroupsOp->emitOpError(
          "unhandled scf.forall op that doesnt have a mapping of "
          "`[#iree_linalg_ext.split_reduction_mapping]`");
    }
    if (failed(IREE::LinalgExt::SplitReductionMappingAttr::verifyAttrList(
            rewriter.getContext(), forallOp.getLoc(), mapping->getValue()))) {
      return failure();
    }
    defaultCountOp =
        IREE::TensorExt::DispatchWorkgroupCountSplitReductionModifierOp::create(
            rewriter, loc, defaultCountOp->getResults(), block->getArguments());
  }
  IREE::Flow::ReturnOp::create(rewriter, loc, defaultCountOp->getResults());

  // Update the `workgroupsOp` region.
  rewriter.modifyOpInPlace(workgroupsOp, [&]() {
    // Update the workload of the op.
    workgroupsOp.getWorkloadMutable().assign(workload);

    // Annotate the values captures as workload with their position in the
    // workload list.
    Region &body = workgroupsOp.getWorkgroupBody();
    if (body.empty()) {
      return;
    }
    rewriter.setInsertionPointToStart(&body.front());
    int ordinalNumber = 0;
    for (auto [index, operand] : llvm::enumerate(workgroupsOp.getArguments())) {
      if (!isa<IndexType>(operand.getType())) {
        continue;
      }
      BlockArgument arg = workgroupsOp.getInputBlockArgument(index);
      auto ordinalOp = IREE::TensorExt::DispatchWorkloadOrdinalOp::create(
          rewriter, loc, arg, rewriter.getIndexAttr(ordinalNumber++));
      rewriter.replaceAllUsesExcept(arg, ordinalOp, ordinalOp);
    }
  });
  return success();
}

namespace {
struct MaterializeDefaultWorkgroupCountRegionPass
    : impl::MaterializeDefaultWorkgroupCountRegionPassBase<
          MaterializeDefaultWorkgroupCountRegionPass> {
  using Base::Base;
  void runOnOperation() override;
};
} // namespace

// populates the workgroup count region.
void MaterializeDefaultWorkgroupCountRegionPass::runOnOperation() {
  FunctionOpInterface funcOp = getOperation();
  TensorDimTrackingRewriter rewriter(funcOp);

  // Populate the workgroup_count region of flow.dispatch.workgroups operation
  // that dont already have a region
  auto walkResult = funcOp.walk(
      [&](IREE::Flow::DispatchWorkgroupsOp workgroupsOp) -> WalkResult {
        if (failed(createDefaultWorkgroupCountRegion(rewriter, workgroupsOp))) {
          return WalkResult::interrupt();
        }
        return WalkResult::advance();
      });
  if (walkResult.wasInterrupted()) {
    return signalPassFailure();
  }
}

} // namespace mlir::iree_compiler::DispatchCreation
