//===- RaisingTransformOps.cpp - Definition of raising transform extension ===//
//
// 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
//
//===----------------------------------------------------------------------===//

#include "src/enzyme_ad/jax/TransformOps/RaisingTransformOps.h"

#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/Dialect/Transform/IR/TransformDialect.h"
#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/PatternMatch.h"

#define GET_OP_CLASSES
#include "src/enzyme_ad/jax/TransformOps/RaisingTransformOps.cpp.inc"
#include "src/enzyme_ad/jax/TransformOps/RaisingTransformOpsImpl.cpp.inc"

using namespace mlir;

namespace mlir {
namespace transform {

// The iteration count comes out in the induction variable's type and is
// multiplied by a step in the iter arg's, which need not be the same type.
// Narrowing the count loses nothing that was ever read: the arg's arithmetic
// wraps, so only the low bits of the count reach the result either way.
// Widening it loses nothing either -- a count is not negative and it fit in
// the narrower type it was counted in.
static Value castToType(OpBuilder &builder, Location loc, Value v, Type ty) {
  if (v.getType() == ty)
    return v;
  if (isa<IndexType>(v.getType()) || isa<IndexType>(ty))
    return arith::IndexCastOp::create(builder, loc, ty, v);
  unsigned from = cast<IntegerType>(v.getType()).getWidth();
  unsigned to = cast<IntegerType>(ty).getWidth();
  if (from > to)
    return arith::TruncIOp::create(builder, loc, ty, v);
  return arith::ExtSIOp::create(builder, loc, ty, v);
}

LogicalResult RemoveIVs::matchAndRewrite(scf::ForOp forOp,
                                         PatternRewriter &rewriter) const {
  if (!forOp.getRegion().hasOneBlock())
    return failure();
  unsigned numIterArgs = forOp.getNumRegionIterArgs();
  auto loc = forOp->getLoc();
  bool changed = false;
  llvm::SetVector<unsigned> removed;
  llvm::MapVector<unsigned, Value> steps;
  auto yield = cast<scf::YieldOp>(forOp.getBody()->getTerminator());

  for (unsigned i = 0; i < numIterArgs; i++) {
    auto ba = forOp.getRegionIterArgs()[i];
    auto init = forOp.getInits()[i];
    auto next = yield->getOperand(i);

    auto increment = next.getDefiningOp<arith::AddIOp>();
    if (!increment)
      continue;

    // Only ba + step is an induction. An addi of two unrelated values (say,
    // an outer accumulator plus this loop's IV) also has a final value, but
    // it is not init + count * step.
    Value step = nullptr;
    if (increment.getLhs() == ba) {
      step = increment.getRhs();
    } else if (increment.getRhs() == ba) {
      step = increment.getLhs();
    }
    if (!step)
      continue;

    // If it dominates the loop entry
    if (!step.getParentRegion()->isProperAncestor(&forOp.getRegion()))
      continue;

    rewriter.setInsertionPointToStart(forOp.getBody());
    Value iterNum = arith::SubIOp::create(
        rewriter, loc, forOp.getInductionVar(), forOp.getLowerBound());
    iterNum = arith::DivSIOp::create(rewriter, loc, iterNum, forOp.getStep());
    iterNum = castToType(rewriter, loc, iterNum, step.getType());

    Value replacementIV = arith::MulIOp::create(rewriter, loc, iterNum, step);
    replacementIV = arith::AddIOp::create(rewriter, loc, replacementIV, init);

    rewriter.replaceAllUsesWith(ba, replacementIV);

    removed.insert(i);
    steps.insert({i, step});
    changed = true;
  }

  if (!changed)
    return failure();

  SmallVector<Value> newInits;
  for (unsigned i = 0; i < numIterArgs; i++)
    if (!removed.contains(i))
      newInits.push_back(forOp.getInits()[i]);

  rewriter.setInsertionPoint(forOp);
  auto newForOp =
      scf::ForOp::create(rewriter, loc, forOp.getLowerBound(),
                         forOp.getUpperBound(), forOp.getStep(), newInits);
  if (!newForOp.getRegion().empty())
    newForOp.getRegion().front().erase();
  assert(newForOp.getRegion().empty());
  rewriter.inlineRegionBefore(forOp.getRegion(), newForOp.getRegion(),
                              newForOp.getRegion().begin());

  SmallVector<Value> newYields;
  for (unsigned i = 0; i < numIterArgs; i++)
    if (!removed.contains(i))
      newYields.push_back(yield->getOperand(i));

  rewriter.setInsertionPoint(yield);
  rewriter.replaceOpWithNewOp<scf::YieldOp>(yield, newYields);

  llvm::BitVector toDelete(numIterArgs + 1);
  for (unsigned i = 0; i < numIterArgs; i++)
    if (removed.contains(i))
      toDelete[i + 1] = true;
  newForOp.getBody()->eraseArguments(toDelete);

  rewriter.setInsertionPoint(newForOp);
  unsigned curNewRes = 0;
  for (unsigned i = 0; i < numIterArgs; i++) {
    auto result = forOp->getResult(i);
    if (removed.contains(i)) {
      if (result.use_empty())
        continue;

      rewriter.setInsertionPointAfter(forOp.getOperation());
      Value iterNum = arith::SubIOp::create(
          rewriter, loc, forOp.getUpperBound(), forOp.getLowerBound());
      iterNum = arith::DivSIOp::create(rewriter, loc, iterNum, forOp.getStep());
      iterNum = castToType(rewriter, loc, iterNum, steps[i].getType());

      Value afterLoop = arith::MulIOp::create(rewriter, loc, iterNum, steps[i]);
      afterLoop =
          arith::AddIOp::create(rewriter, loc, afterLoop, forOp.getInits()[i]);

      rewriter.replaceAllUsesWith(result, afterLoop);
    } else {
      rewriter.replaceAllUsesWith(result, newForOp->getResult(curNewRes++));
    }
  }

  rewriter.eraseOp(forOp);

  return success();
}

static inline void clearBlock(mlir::Block *block,
                              mlir::RewriterBase &rewriter) {
  for (auto &op : llvm::make_early_inc_range(llvm::reverse(*block))) {
    assert(op.use_empty() && "expected 'op' to have no uses");
    rewriter.eraseOp(&op);
  }
}

static mlir::Value createConstantInt(RewriterBase &rewriter, Location loc,
                                     Type ty, int64_t v) {
  if (ty.isIndex())
    return arith::ConstantIndexOp::create(rewriter, loc, v);
  else
    return arith::ConstantIntOp::create(rewriter, loc, ty, v);
}

static std::optional<int64_t> getConstant(Operation *op) {
  if (auto cst = dyn_cast_or_null<arith::ConstantIntOp>(op)) {
    return cst.value();
  } else if (auto cst = dyn_cast_or_null<arith::ConstantIndexOp>(op)) {
    return cst.value();
  } else if (auto cst = dyn_cast_or_null<LLVM::ConstantOp>(op)) {
    if (auto intAttr = dyn_cast<IntegerAttr>(cst.getValue()))
      return intAttr.getValue().getSExtValue();
  }
  return {};
}

static std::optional<int64_t> getConstant(Value v) {
  Operation *op = v.getDefiningOp();
  if (op)
    return getConstant(op);
  return {};
}

/// Returns `true` if the loop has a form expected by interchange patterns.
static bool isNormalized(scf::ForOp op) {
  auto lb = getConstant(op.getLowerBound());
  auto step = getConstant(op.getStep());
  if (!lb || !step)
    return false;
  return *lb == 0 && *step == 1;
}

#define DEBUG_TYPE "normalize-loop"
#define DBGS llvm::dbgs

LogicalResult NormalizeLoop::matchAndRewrite(scf::ForOp op,
                                             PatternRewriter &rewriter) const {
  using namespace arith;
  if (isNormalized(op) ||
      !isa<scf::ParallelOp, affine::AffineParallelOp>(op->getParentOp())) {
    LLVM_DEBUG(DBGS() << "[normalize-loop] loop already normalized\n");
    return failure();
  }

  rewriter.setInsertionPoint(op);
  Value zero = createConstantInt(rewriter, op.getLoc(),
                                 op.getInductionVar().getType(), 0);
  Value one = createConstantInt(rewriter, op.getLoc(),
                                op.getInductionVar().getType(), 1);

  Value difference = SubIOp::create(rewriter, op.getLoc(), op.getUpperBound(),
                                    op.getLowerBound());
  Value tripCount = AddIOp::create(
      rewriter, op.getLoc(),
      DivUIOp::create(rewriter, op.getLoc(),
                      SubIOp::create(rewriter, op.getLoc(), difference, one),
                      op.getStep()),
      one);
  auto newForOp = scf::ForOp::create(rewriter, op.getLoc(), zero, tripCount,
                                     one, op.getInits());
  clearBlock(newForOp.getBody(), rewriter);
  rewriter.setInsertionPointToStart(newForOp.getBody());
  Value scaled = MulIOp::create(rewriter, op.getLoc(),
                                newForOp.getInductionVar(), op.getStep());
  Value iv = AddIOp::create(rewriter, op.getLoc(), op.getLowerBound(), scaled);
  SmallVector<Value> newArgs(newForOp.getRegion().args_begin(),
                             newForOp.getRegion().args_end());
  newArgs[0] = iv;
  rewriter.inlineBlockBefore(op.getBody(), newForOp.getBody(),
                             newForOp.getBody()->end(), newArgs);
  rewriter.replaceOp(op, newForOp->getResults());
  return success();
}

} // namespace transform
} // namespace mlir
#include "src/enzyme_ad/jax/TransformOps/RaisingTransformPatterns.cpp.inc"

namespace {
class RaisingTransformExtension
    : public transform::TransformDialectExtension<RaisingTransformExtension> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(RaisingTransformExtension)
  using Base::Base;

  void init() {
    registerTransformOps<
#define GET_OP_LIST
#include "src/enzyme_ad/jax/TransformOps/RaisingTransformOps.cpp.inc"
        >();
  }
};
} // namespace

void mlir::enzyme::registerRaisingTransformExtension(
    DialectRegistry &registry) {
  registry.addExtensions<RaisingTransformExtension>();
}
