// Copyright 2025 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/Passes.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/Interfaces/LoopLikeInterface.h"

namespace mlir::iree_compiler {

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

namespace {
struct ResolveSwizzleHintsPass final
    : impl::ResolveSwizzleHintsPassBase<ResolveSwizzleHintsPass> {
  using Base::Base;
  void runOnOperation() override;
};
} // namespace

static Value createOrFoldNewStaticAdd(RewriterBase &rewriter, Value v,
                                      int64_t offset) {
  // Early exit for the common offset = 0 case.
  if (offset == 0) {
    return v;
  }

  if (auto add = v.getDefiningOp<arith::AddIOp>()) {
    llvm::APInt constant;
    if (matchPattern(add.getRhs(), m_ConstantInt(&constant))) {
      Value combined = arith::ConstantIndexOp::create(
          rewriter, add.getLoc(), offset + constant.getSExtValue());
      return arith::AddIOp::create(rewriter, add.getLoc(), add.getLhs(),
                                   combined, add.getOverflowFlags());
    }
  }
  Value offsetVal =
      arith::ConstantIndexOp::create(rewriter, v.getLoc(), offset);
  return arith::AddIOp::create(rewriter, v.getLoc(), v, offsetVal);
}

/// Swizzles vector.load(iree_codegen.swizzle_hint, offset). The
/// SwizzleInterfaceAttr exposes two methods:
///   1. getAccessElementCount -> int64_t
///        - Gives the number of contiguous elements in the swizzling pattern.
///   2. swizzleOffset(src: memref<N x !eltype>, offset: index) -> index
///        - Swizzles the |offset| into |src|, returning the new offset.
///
/// For a 1-d load of type `vector<C x !eltype>`, the load is unrolled into
/// loads of size `k = getAccessElementCount()` and individually swizzled.
///
/// For example, if `C = 16` and `k = 4`, this produces:
///
/// %0 = vector.load %src[swizzleOffset(%src, %offset)] : vector<4>
/// %1 = vector.load %src[swizzleOffset(%src, %offset + 4)] : vector<4>
/// %2 = vector.load %src[swizzleOffset(%src, %offset + 8)] : vector<4>
/// %3 = vector.load %src[swizzleOffset(%src, %offset + 12)] : vector<4>
/// %load = concat[%0, %1, %2, %3] : vector<16>
static void swizzleLoad(RewriterBase &rewriter, vector::LoadOp load,
                        IREE::Codegen::SwizzleHintOp hintOp) {
  Location hintLoc = hintOp.getLoc();
  int64_t accessWidth = hintOp.getSwizzle().getAccessElementCount();
  VectorType type = load.getVectorType();
  int64_t loadWidth = type.getShape()[0];
  Value memrefOffset = load.getIndices()[0];
  VectorType swizzledLoadType =
      VectorType::get({accessWidth}, type.getElementType());

  // ~ vector.undef, overwritten by unrolling.
  Value replacement = arith::ConstantOp::create(rewriter, hintLoc, type,
                                                rewriter.getZeroAttr(type));

  // Load type = vector<C>, k = accessWidth
  // i = 0 -> C += k is the offset into the vector of a contiguous group of
  // swizzled elements.
  for (int64_t i = 0; i < loadWidth; i += accessWidth) {
    Value newBaseOffset = createOrFoldNewStaticAdd(rewriter, memrefOffset, i);
    Value newOffset = getValueOrCreateConstantIndexOp(
        rewriter, hintLoc,
        hintOp.getSwizzle().swizzleOffset(rewriter, hintOp.getLoc(),
                                          newBaseOffset, hintOp.getOperand()));
    auto subLoad = vector::LoadOp::create(
        rewriter, load.getLoc(), swizzledLoadType, load.getBase(), newOffset);

    replacement = vector::InsertStridedSliceOp::create(
        rewriter, load.getLoc(), subLoad, replacement, ArrayRef<int64_t>{i},
        ArrayRef<int64_t>{1});
  }
  rewriter.replaceOp(load, replacement);
}

/// Swizzles vector.store(iree_codegen.swizzle_hint, offset).
///
/// For a 1-d store of type `vector<C x !eltype>`, the store is unrolled into
/// stores of size `k = getAccessElementCount()` and individually swizzled.
///
/// For example, if `C = 16` and `k = 4`, this produces:
///
/// %0, %1, %2, %3 = split[%value_to_store] : vector<16>
/// vector.store %0, %src[swizzleOffset(%src, %offset)] : vector<4>
/// vector.store %1, %src[swizzleOffset(%src, %offset + 4)] : vector<4>
/// vector.store %2, %src[swizzleOffset(%src, %offset + 8)] : vector<4>
/// vector.store %3, %src[swizzleOffset(%src, %offset + 12)] : vector<4>
static void swizzleStore(RewriterBase &rewriter, vector::StoreOp store,
                         IREE::Codegen::SwizzleHintOp hintOp) {
  Location hintLoc = hintOp.getLoc();
  int64_t accessWidth = hintOp.getSwizzle().getAccessElementCount();
  VectorType type = store.getVectorType();
  int64_t storeWidth = type.getShape()[0];
  Value memrefOffset = store.getIndices()[0];

  // Store type = vector<C>, k = accessWidth
  // i = 0 -> C += k is the offset into the vector of a contiguous group of
  // swizzled elements.
  for (int64_t i = 0; i < storeWidth; i += accessWidth) {
    Value subVec = vector::ExtractStridedSliceOp::create(
        rewriter, store.getLoc(), store.getValueToStore(), ArrayRef<int64_t>{i},
        ArrayRef<int64_t>{accessWidth}, ArrayRef<int64_t>{1});
    Value newBaseOffset = createOrFoldNewStaticAdd(rewriter, memrefOffset, i);

    Value newOffset = getValueOrCreateConstantIndexOp(
        rewriter, hintLoc,
        hintOp.getSwizzle().swizzleOffset(rewriter, hintOp.getLoc(),
                                          newBaseOffset, hintOp.getOperand()));
    vector::StoreOp::create(rewriter, store.getLoc(), subVec, store.getBase(),
                            newOffset);
  }
  rewriter.eraseOp(store);
}

static LogicalResult
verifyFlatContiguousSwizzleHintOp(IREE::Codegen::SwizzleHintOp hintOp) {
  auto memrefType = cast<MemRefType>(hintOp.getOperand().getType());
  // Swizzle hints require flat (rank 1) memrefs.
  // For rank 1, allow dynamic memrefs or static contiguous row-major memrefs.
  if (memrefType.getRank() != 1 ||
      (memrefType.hasStaticShape() &&
       !memref::isStaticShapeAndContiguousRowMajor(memrefType))) {
    hintOp.emitError()
        << "swizzle hint operand must be a contiguous flat memref, got "
        << hintOp.getOperand().getType();
    return failure();
  }
  return success();
}

/// Traces a memref value backward through defining ops and loop
/// iter_args/results to find the root memref.alloc.
static memref::AllocOp traceToAllocation(Value val) {
  DenseSet<Value> visited;
  SmallVector<Value> worklist = {val};
  while (!worklist.empty()) {
    Value current = worklist.pop_back_val();
    if (!visited.insert(current).second) {
      continue;
    }
    if (auto blockArg = dyn_cast<BlockArgument>(current)) {
      auto loopOp =
          dyn_cast<LoopLikeOpInterface>(blockArg.getOwner()->getParentOp());
      if (!loopOp) {
        continue;
      }
      if (OpOperand *init = loopOp.getTiedLoopInit(blockArg)) {
        worklist.push_back(init->get());
      }
    } else if (auto allocOp = current.getDefiningOp<memref::AllocOp>()) {
      return allocOp;
    } else if (auto loopOp = current.getDefiningOp<LoopLikeOpInterface>()) {
      if (OpOperand *init = loopOp.getTiedLoopInit(cast<OpResult>(current))) {
        worklist.push_back(init->get());
      }
    } else {
      for (Value operand : current.getDefiningOp()->getOperands()) {
        if (isa<MemRefType>(operand.getType())) {
          worklist.push_back(operand);
        }
      }
    }
  }
  return nullptr;
}

struct SwizzleState {
  bool resolvable = true;
  bool hasGatherToLds = false;
};

/// Verifies the users of a swizzle hint and returns its swizzle state.
/// Emits a diagnostic for any unsupported user.
static SwizzleState verifyHintUsers(IREE::Codegen::SwizzleHintOp hintOp) {
  SwizzleState state;
  int64_t accessWidth = hintOp.getSwizzle().getAccessElementCount();
  for (Operation *user : hintOp->getUsers()) {
    if (auto load = dyn_cast<vector::LoadOp>(user)) {
      VectorType loadType = load.getVectorType();
      // Guard on zero rank loads and loads not divisible by the access width.
      if (loadType.getRank() != 1 ||
          loadType.getShape()[0] % accessWidth != 0) {
        state.resolvable = false;
      }
      continue;
    }
    if (auto store = dyn_cast<vector::StoreOp>(user)) {
      VectorType storeType = store.getVectorType();
      // Guard on zero rank stores and stores not divisible by the access width.
      if (storeType.getRank() != 1 ||
          storeType.getShape()[0] % accessWidth != 0) {
        state.resolvable = false;
      }
      continue;
    }
    if (isa<amdgpu::GatherToLDSOp>(user)) {
      state.hasGatherToLds = true;
      continue;
    }
    hintOp.emitError() << "unsupported SwizzleHintOp user: " << *user;
    state.resolvable = false;
  }
  return state;
}

/// Rewrites all load/store users of a resolved swizzle hint with swizzled
/// offsets. Only call after verifyHintUsers returned a resolvable state.
static void resolveHintOp(RewriterBase &rewriter,
                          IREE::Codegen::SwizzleHintOp hintOp) {
  for (Operation *user : llvm::make_early_inc_range(hintOp->getUsers())) {
    if (auto load = dyn_cast<vector::LoadOp>(user)) {
      rewriter.setInsertionPoint(load);
      swizzleLoad(rewriter, load, hintOp);
    } else if (auto store = dyn_cast<vector::StoreOp>(user)) {
      rewriter.setInsertionPoint(store);
      swizzleStore(rewriter, store, hintOp);
    }
  }
}

void ResolveSwizzleHintsPass::runOnOperation() {
  FunctionOpInterface funcOp = getOperation();

  // Collect all hint ops.
  SmallVector<IREE::Codegen::SwizzleHintOp> hintOps;
  funcOp.walk(
      [&](IREE::Codegen::SwizzleHintOp hint) { hintOps.push_back(hint); });

  // Verify resolvability of each hint and group results by alloc, since hints
  // sharing an alloc must be resolved or dropped together.
  DenseMap<Value, SwizzleState> allocToSwizzleState;
  DenseMap<Operation *, Value> swizzleOpToAlloc;

  for (IREE::Codegen::SwizzleHintOp hintOp : hintOps) {
    if (failed(verifyFlatContiguousSwizzleHintOp(hintOp))) {
      return signalPassFailure();
    }
    // Hints normally trace to a memref.alloc. BlockArgument only occurs in lit
    // tests.
    Value alloc;
    if (memref::AllocOp allocOp = traceToAllocation(hintOp.getOperand())) {
      alloc = allocOp.getResult();
    } else {
      assert(isa<BlockArgument>(hintOp.getOperand()) &&
             "hint operand without alloc should be a function argument");
      alloc = hintOp.getOperand();
    }
    swizzleOpToAlloc[hintOp] = alloc;

    SwizzleState hintState = verifyHintUsers(hintOp);
    SwizzleState &state = allocToSwizzleState[alloc];
    state.resolvable &= hintState.resolvable;
    state.hasGatherToLds |= hintState.hasGatherToLds;
  }

  // gather_to_lds is a special case: its swizzle is already resolved before
  // this pass runs, so we cannot silently drop it here.
  for (auto &[alloc, state] : allocToSwizzleState) {
    memref::AllocOp allocOp = alloc.getDefiningOp<memref::AllocOp>();
    if (!allocOp) {
      continue;
    }
    if (!state.resolvable && state.hasGatherToLds) {
      allocOp.emitError()
          << "alloc has gather_to_lds users; all swizzle hints must resolve";
      return signalPassFailure();
    }
  }

  // Resolve hints.
  IRRewriter rewriter(funcOp->getContext());
  for (IREE::Codegen::SwizzleHintOp hintOp : hintOps) {
    Value alloc = swizzleOpToAlloc[hintOp];
    if (allocToSwizzleState[alloc].resolvable) {
      resolveHintOp(rewriter, hintOp);
    }
  }

  // Drop all hints.
  for (auto hintOp : hintOps) {
    rewriter.replaceOp(hintOp, hintOp.getOperand());
  }
}

} // namespace mlir::iree_compiler
