//===- OpenACCUtilsLoop.cpp - OpenACC Loop Utilities ----------------------===//
//
// Part of the LLVM Project, 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
//
//===----------------------------------------------------------------------===//
//
// This file contains utility functions for converting OpenACC loops to SCF.
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/OpenACC/OpenACCUtilsLoop.h"

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/OpenACC/OpenACC.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/IR/IRMapping.h"

using namespace mlir;

namespace {

/// Calculate trip count for a loop: (ub - lb + step) / step
/// If inclusiveUpperbound is false, subtracts 1 from ub first.
static Value calculateTripCount(OpBuilder &b, Location loc, Value lb, Value ub,
                                Value step, bool inclusiveUpperbound) {
  Type type = b.getIndexType();

  // Convert original loop arguments to index type
  lb = getValueOrCreateCastToIndexLike(b, loc, type, lb);
  ub = getValueOrCreateCastToIndexLike(b, loc, type, ub);
  step = getValueOrCreateCastToIndexLike(b, loc, type, step);

  if (!inclusiveUpperbound) {
    Value one = arith::ConstantIndexOp::create(b, loc, 1);
    ub = b.createOrFold<arith::SubIOp>(loc, ub, one,
                                       arith::IntegerOverflowFlags::nsw);
  }

  Value sub = b.createOrFold<arith::SubIOp>(loc, ub, lb,
                                            arith::IntegerOverflowFlags::nsw);
  Value add = b.createOrFold<arith::AddIOp>(loc, sub, step,
                                            arith::IntegerOverflowFlags::nsw);
  return b.createOrFold<arith::DivSIOp>(loc, add, step);
}

/// Get exclusive upper bound from acc.loop (add 1 if inclusive).
/// The result is always in index type.
static Value getExclusiveUpperBoundAsIndex(acc::LoopOp loopOp, size_t ivPos,
                                           OpBuilder &b) {
  bool isInclusive = false;
  if (loopOp.getInclusiveUpperbound().has_value())
    isInclusive = loopOp.getInclusiveUpperboundAttr().asArrayRef()[ivPos];

  Value origUB = loopOp.getUpperbound()[ivPos];
  Location loc = origUB.getLoc();
  Type indexType = b.getIndexType();

  // Cast to index first, then add if inclusive
  Value ub = getValueOrCreateCastToIndexLike(b, loc, indexType, origUB);
  if (isInclusive) {
    Value one = arith::ConstantIndexOp::create(b, loc, 1);
    ub = b.createOrFold<arith::AddIOp>(loc, ub, one,
                                       arith::IntegerOverflowFlags::nsw);
  }
  return ub;
}

/// Handle differing types between SCF (index) and ACC loops.
/// Creates casts from the new SCF IVs to the original ACC IV types and updates
/// the mapping. The newIVs should correspond 1:1 with the ACC loop's IVs.
static void mapACCLoopIVsToSCFIVs(acc::LoopOp accLoop, ValueRange newIVs,
                                  OpBuilder &b, IRMapping &mapping) {
  for (auto [origIV, newIV] :
       llvm::zip(accLoop.getBody().getArguments(), newIVs)) {
    Value replacementIV = getValueOrCreateCastToIndexLike(
        b, accLoop->getLoc(), origIV.getType(), newIV);
    mapping.map(origIV, replacementIV);
  }
}

/// Normalize IV uses after converting to normalized loop form.
/// For normalized loops (lb=0, step=1), we need to denormalize the IV:
/// original_iv = new_iv * orig_step + orig_lb
static void normalizeIVUses(OpBuilder &b, Location loc, Value iv, Value origLB,
                            Value origStep) {
  Type indexType = b.getIndexType();
  Value lb = getValueOrCreateCastToIndexLike(b, loc, indexType, origLB);
  Value step = getValueOrCreateCastToIndexLike(b, loc, indexType, origStep);

  // new_iv * step + lb
  Value scaled =
      arith::MulIOp::create(b, loc, iv, step, arith::IntegerOverflowFlags::nsw);
  Value denormalized = arith::AddIOp::create(b, loc, scaled, lb,
                                             arith::IntegerOverflowFlags::nsw);

  // Replace uses of iv with denormalized value, except for the ops that
  // compute the denormalized value itself (muli and addi)
  llvm::SmallPtrSet<Operation *, 2> exceptions;
  exceptions.insert(scaled.getDefiningOp());
  exceptions.insert(denormalized.getDefiningOp());
  iv.replaceAllUsesExcept(denormalized, exceptions);
}

/// Clone an ACC region into a destination block, handling the ACC terminators.
/// Returns the insertion point after the cloned operations.
static Block::iterator cloneACCRegionInto(Region *src, Block *dest,
                                          Block::iterator insertionPoint,
                                          IRMapping &mapping,
                                          RewriterBase &rewriter) {
  assert(src->hasOneBlock() && "expected single-block region");

  Region *insertRegion = dest->getParent();
  Block *postInsertBlock = rewriter.splitBlock(dest, insertionPoint);
  rewriter.cloneRegionBefore(*src, *insertRegion,
                             postInsertBlock->getIterator(), mapping);

  auto lastNewBlock = std::prev(postInsertBlock->getIterator());

  Block::iterator newInsertionPoint;
  Operation *terminator = lastNewBlock->getTerminator();

  if (auto yieldOp = dyn_cast<acc::YieldOp>(terminator)) {
    newInsertionPoint = std::prev(yieldOp->getIterator());
    rewriter.eraseOp(yieldOp);
  } else if (auto terminatorOp = dyn_cast<acc::TerminatorOp>(terminator)) {
    newInsertionPoint = std::prev(terminatorOp->getIterator());
    rewriter.eraseOp(terminatorOp);
  } else {
    llvm_unreachable("unexpected terminator in ACC region");
  }

  // Merge last block with the postInsertBlock
  rewriter.mergeBlocks(postInsertBlock, &*lastNewBlock);

  // Merge first block with original dest block
  Block *firstNewBlock = &*std::next(dest->getIterator());
  rewriter.mergeBlocks(firstNewBlock, dest);

  return newInsertionPoint;
}

/// Wrap a multi-block region with scf.execute_region.
static scf::ExecuteRegionOp
wrapMultiBlockRegionWithSCFExecuteRegion(Region &region, IRMapping &mapping,
                                         Location loc, RewriterBase &rewriter) {
  auto exeRegionOp = scf::ExecuteRegionOp::create(rewriter, loc, TypeRange{});

  rewriter.cloneRegionBefore(region, exeRegionOp.getRegion(),
                             exeRegionOp.getRegion().end(), mapping);

  // Find and replace the ACC terminator with scf.yield
  Operation *terminator = exeRegionOp.getRegion().back().getTerminator();
  if (auto yieldOp = dyn_cast<acc::YieldOp>(terminator)) {
    if (yieldOp.getNumOperands() > 0) {
      region.getParentOp()->emitError(
          "acc.loop with results not yet supported");
      return nullptr;
    }
  } else if (!isa<acc::TerminatorOp>(terminator)) {
    llvm_unreachable("unexpected terminator in ACC region");
  }

  rewriter.eraseOp(terminator);
  rewriter.setInsertionPointToEnd(&exeRegionOp.getRegion().back());
  scf::YieldOp::create(rewriter, loc);
  return exeRegionOp;
}

} // namespace

namespace mlir {
namespace acc {

scf::ForOp convertACCLoopToSCFFor(LoopOp loopOp, RewriterBase &rewriter,
                                  bool enableCollapse) {
  assert(!loopOp.getUnstructured() &&
         "use convertUnstructuredACCLoopToSCFExecuteRegion for unstructured "
         "loops");

  Location loc = loopOp->getLoc();
  Type indexType = rewriter.getIndexType();

  // Create nested scf.for loops and build IR mapping for IVs
  IRMapping mapping;
  SmallVector<scf::ForOp> forOps;

  // Save the original insertion point
  OpBuilder::InsertionGuard guard(rewriter);
  rewriter.setInsertionPoint(loopOp);

  // First, compute ALL loop bounds at the current insertion point (before
  // any ForOp). This ensures all bounds are defined in the outer scope,
  // which is required for coalesceLoops to work correctly.
  SmallVector<Value> lowerBounds, upperBounds, steps;
  for (BlockArgument iv : loopOp.getBody().getArguments()) {
    size_t idx = iv.getArgNumber();
    Value newLowerBound = getValueOrCreateCastToIndexLike(
        rewriter, loc, indexType, loopOp.getLowerbound()[idx]);
    Value newUpperBound = getExclusiveUpperBoundAsIndex(loopOp, idx, rewriter);
    Value newStep = getValueOrCreateCastToIndexLike(rewriter, loc, indexType,
                                                    loopOp.getStep()[idx]);
    lowerBounds.push_back(newLowerBound);
    upperBounds.push_back(newUpperBound);
    steps.push_back(newStep);
  }

  // Now create the nested ForOps using the pre-computed bounds
  for (BlockArgument iv : loopOp.getBody().getArguments()) {
    size_t idx = iv.getArgNumber();

    // For nested loops, insert inside the previous loop's body
    if (idx > 0)
      rewriter.setInsertionPointToStart(forOps.back().getBody());

    scf::ForOp forOp = scf::ForOp::create(rewriter, loc, lowerBounds[idx],
                                          upperBounds[idx], steps[idx]);
    forOps.push_back(forOp);
    mapping.map(iv, forOp.getInductionVar());
  }

  // Set insertion point inside the innermost loop for IV casts and body cloning
  rewriter.setInsertionPointToStart(forOps.back().getBody());

  // Handle IV type conversion (index -> original type)
  SmallVector<Value> scfIVs;
  for (scf::ForOp forOp : forOps)
    scfIVs.push_back(forOp.getInductionVar());
  mapACCLoopIVsToSCFIVs(loopOp, scfIVs, rewriter, mapping);

  // Clone the loop body into the innermost scf.for
  cloneACCRegionInto(&loopOp.getRegion(), forOps.back().getBody(),
                     rewriter.getInsertionPoint(), mapping, rewriter);

  // Optionally collapse nested loops
  if (enableCollapse && forOps.size() > 1)
    if (failed(coalesceLoops(rewriter, forOps)))
      loopOp.emitError("failed to collapse acc.loop");

  return forOps.front();
}

scf::ParallelOp convertACCLoopToSCFParallel(LoopOp loopOp,
                                            RewriterBase &rewriter) {
  assert(!loopOp.getUnstructured() &&
         "use convertUnstructuredACCLoopToSCFExecuteRegion for unstructured "
         "loops");
  assert(
      rewriter.getInsertionBlock() &&
      !loopOp->isProperAncestor(rewriter.getInsertionBlock()->getParentOp()) &&
      "builder insertion point must not be inside the loop being converted");

  Location loc = loopOp->getLoc();

  SmallVector<Value> lowerBounds, upperBounds, steps;

  // Normalize all loops: lb=0, step=1, ub=tripCount
  Value lb = arith::ConstantIndexOp::create(rewriter, loc, 0);
  Value step = arith::ConstantIndexOp::create(rewriter, loc, 1);

  for (auto [idx, iv] : llvm::enumerate(loopOp.getBody().getArguments())) {
    bool inclusiveUpperbound = false;
    if (loopOp.getInclusiveUpperbound().has_value())
      inclusiveUpperbound = loopOp.getInclusiveUpperbound().value()[idx];

    Value ub = calculateTripCount(rewriter, loc, loopOp.getLowerbound()[idx],
                                  loopOp.getUpperbound()[idx],
                                  loopOp.getStep()[idx], inclusiveUpperbound);

    lowerBounds.push_back(lb);
    upperBounds.push_back(ub);
    steps.push_back(step);
  }

  auto parallelOp =
      scf::ParallelOp::create(rewriter, loc, lowerBounds, upperBounds, steps);

  // Create IV type conversions
  IRMapping mapping;
  rewriter.setInsertionPointToStart(parallelOp.getBody());
  mapACCLoopIVsToSCFIVs(loopOp, parallelOp.getInductionVars(), rewriter,
                        mapping);

  if (!loopOp.getRegion().hasOneBlock()) {
    auto exeRegion = wrapMultiBlockRegionWithSCFExecuteRegion(
        loopOp.getRegion(), mapping, loc, rewriter);
    if (!exeRegion) {
      rewriter.eraseOp(parallelOp);
      return nullptr;
    }
  } else {
    cloneACCRegionInto(&loopOp.getRegion(), parallelOp.getBody(),
                       rewriter.getInsertionPoint(), mapping, rewriter);
  }

  // Denormalize IV uses
  rewriter.setInsertionPointToStart(parallelOp.getBody());
  for (auto [idx, iv] : llvm::enumerate(parallelOp.getBody()->getArguments()))
    if (!iv.use_empty())
      normalizeIVUses(rewriter, loc, iv, loopOp.getLowerbound()[idx],
                      loopOp.getStep()[idx]);

  return parallelOp;
}

scf::ExecuteRegionOp
convertUnstructuredACCLoopToSCFExecuteRegion(LoopOp loopOp,
                                             RewriterBase &rewriter) {
  assert(loopOp.getUnstructured() &&
         "use convertACCLoopToSCFFor for structured loops");
  assert(
      rewriter.getInsertionBlock() &&
      !loopOp->isProperAncestor(rewriter.getInsertionBlock()->getParentOp()) &&
      "builder insertion point must not be inside the loop being converted");

  IRMapping mapping;
  return wrapMultiBlockRegionWithSCFExecuteRegion(loopOp.getRegion(), mapping,
                                                  loopOp->getLoc(), rewriter);
}

} // namespace acc
} // namespace mlir